@@ -8,6 +8,7 @@ import grpc
88from grpc.experimental import aio
99import math
1010import pytest
11+ from proto.marshal.rules.dates import DurationRule, TimestampRule
1112
1213{# Import the service itself as well as every proto module that it imports. -#}
1314{% filter sort_lines -%}
@@ -574,7 +575,13 @@ def test_{{ method.name|snake_case }}_flattened():
574575 assert len(call.mock_calls) == 1
575576 _, args, _ = call.mock_calls[0]
576577 {% for key , field in method .flattened_fields .items () -%}{% - if not field .oneof or field .proto 3_optional %}
578+ {% if field .ident |string () == 'timestamp.Timestamp' -%}
579+ assert TimestampRule().to_proto(args[0].{{ key }}) == {{ field.mock_value }}
580+ {% elif field .ident |string () == 'duration.Duration' -%}
581+ assert DurationRule().to_proto(args[0].{{ key }}) == {{ field.mock_value }}
582+ {% else -%}
577583 assert args[0].{{ key }} == {{ field.mock_value }}
584+ {% endif %}
578585 {% endif %}{% endfor %}
579586 {% - for oneofs in method .flattened_oneof_fields ().values () %}
580587 {% - with field = oneofs [-1] %}
@@ -653,7 +660,13 @@ async def test_{{ method.name|snake_case }}_flattened_async():
653660 assert len(call.mock_calls)
654661 _, args, _ = call.mock_calls[0]
655662 {% for key , field in method .flattened_fields .items () -%}{% - if not field .oneof or field .proto 3_optional %}
663+ {% if field .ident |string () == 'timestamp.Timestamp' -%}
664+ assert TimestampRule().to_proto(args[0].{{ key }}) == {{ field.mock_value }}
665+ {% elif field .ident |string () == 'duration.Duration' -%}
666+ assert DurationRule().to_proto(args[0].{{ key }}) == {{ field.mock_value }}
667+ {% else -%}
656668 assert args[0].{{ key }} == {{ field.mock_value }}
669+ {% endif %}
657670 {% endif %}{% endfor %}
658671 {% - for oneofs in method .flattened_oneof_fields ().values () %}
659672 {% - with field = oneofs [-1] %}
0 commit comments