address review feedback

Author: parthea

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/_shared_macros.j2

@@ -402,7 +402,7 @@ class {{ async_method_name_prefix }}{{ service.name }}RestInterceptor:
                 logging.log(f"Received request: {request}")
                 return request, metadata
 
-    {% if not method.void %}
+    {% if not method.void and not method.is_resumable_upload %}
             {{ async_prefix }}def post_{{ method.name|snake_case }}(self, response):
                 logging.log(f"Received response: {response}")
                 return response
@@ -425,7 +425,7 @@ class {{ async_method_name_prefix }}{{ service.name }}RestInterceptor:
         """
         return request, metadata
 
-    {% if not method.void %}
+    {% if not method.void and not method.is_resumable_upload %}
     {% if not method.server_streaming %}
     {{ async_prefix }}def post_{{ method.name|snake_case }}(self, response: {{method.output.ident}}) -> {{method.output.ident}}:
     {% else %}

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/grpc_asyncio.py.j2

@@ -375,30 +375,32 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
                             )
                         return _error()
                 self._stubs['{{ method.transport_safe_name|snake_case }}'] = _ErrorStub()
-            elif HAS_ASYNC_REST:  # pragma: NO COVER
-                # Lazily instantiate and cache the AsyncRestTransport stub on first
-                # call rather than during __init__ (_prep_wrapped_messages), since
-                # GrpcAsyncIOTransport may be initialized with synchronous
-                # google.auth.credentials.Credentials that AsyncAuthorizedSession
-                # rejects, and gRPC async stubs expose a synchronous __call__
-                # returning an awaitable.
-                class _AsyncRestStub:  # pragma: NO COVER
-                    def __init__(_self):  # pragma: NO COVER
-                        _self._rest_stub = None
-
-                    def __call__(_self, *args, **kwargs):  # pragma: NO COVER
-                        if _self._rest_stub is None:
+            elif HAS_ASYNC_REST:
+                # Defer instantiating Async{{ service.name }}RestTransport until the
+                # resumable upload RPC is called so that initializing
+                # {{ service.grpc_asyncio_transport_name }} with synchronous
+                # google.auth.credentials.Credentials (e.g. via ADC or service
+                # account helpers) for standard gRPC calls does not fail in
+                # __init__ (_prep_wrapped_messages) when AsyncAuthorizedSession
+                # validates for async credentials.
+                transport = self
+                class _AsyncRestStub:
+                    def __init__(self):
+                        self._rest_stub = None
+
+                    def __call__(self, *args, **kwargs):
+                        if self._rest_stub is None:
                             rest_transport = Async{{ service.name }}RestTransport(
-                                host=self._host,
-                                credentials=self._credentials,
-                                client_info=self._client_info,
+                                host=transport._host,
+                                credentials=transport._credentials,
+                                client_info=transport._client_info,
                             )
-                            _self._rest_stub = rest_transport.{{ method.transport_safe_name|snake_case }}
-                        return _self._rest_stub(*args, **kwargs)
+                            self._rest_stub = rest_transport.{{ method.transport_safe_name|snake_case }}
+                        return self._rest_stub(*args, **kwargs)
                 self._stubs['{{ method.transport_safe_name|snake_case }}'] = _AsyncRestStub()
-            else:  # pragma: NO COVER
-                class _UnsupportedStub:  # pragma: NO COVER
-                    def __call__(self, *args, **kwargs):  # pragma: NO COVER
+            else:
+                class _UnsupportedStub:
+                    def __call__(self, *args, **kwargs):
                         async def _unsupported():
                             raise NotImplementedError("Async REST transport is required for async resumable upload methods.")
                         return _unsupported()

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest.py.j2

@@ -277,12 +277,10 @@ class {{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
 
             uri = transcoded_request["uri"]
             upload_url = f"{self._host}{uri}"
-            headers = dict(metadata)
+            headers = {**dict(metadata), **dict((config.headers or {}) if config else {})}
             if config is None:
                 config = resumable_transfer.ResumableUploadConfig(headers=headers)
             else:
-                if config.headers:
-                    headers.update(dict(config.headers))
                 config = dataclasses.replace(config, headers=headers)
 
             return {{ method.client_output.ident }}(

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2

@@ -236,12 +236,10 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
 
             uri = transcoded_request["uri"]
             upload_url = f"{self._host}{uri}"
-            headers = dict(metadata)
+            headers = {**dict(metadata), **dict((config.headers or {}) if config else {})}
             if config is None:
                 config = resumable_transfer.ResumableUploadConfig(headers=headers)
             else:
-                if config.headers:
-                    headers.update(dict(config.headers))
                 config = dataclasses.replace(config, headers=headers)
 
             return {{ method.client_output_async.ident }}(

packages/gapic-generator/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

@@ -1545,6 +1545,50 @@ async def test_{{ service.name|snake_case }}_{{ method.name|snake_case }}_grpc_a
     with pytest.raises(core_exceptions.GoogleAPICallError) as exc_info:
         await transport.{{ method.transport_safe_name|snake_case }}({{ method.input.ident }}())
     assert "operate over REST and cannot be invoked when the transport is initialized with a pre-constructed gRPC channel" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_{{ service.name|snake_case }}_{{ method.name|snake_case }}_grpc_asyncio_rest_stub():
+    with mock.patch.object(transports.grpc_asyncio, "HAS_ASYNC_REST", True), mock.patch.object(
+        transports.grpc_asyncio,
+        "Async{{ service.name }}RestTransport",
+        create=True,
+    ) as mock_rest_transport_cls:
+        mock_rest_transport = mock.MagicMock()
+        expected_session = {{ method.client_output_async.ident }}()
+        mock_rest_transport.{{ method.transport_safe_name|snake_case }} = mock.AsyncMock(
+            return_value=expected_session
+        )
+        mock_rest_transport_cls.return_value = mock_rest_transport
+
+        transport = transports.{{ service.name }}GrpcAsyncIOTransport(
+            credentials=async_anonymous_credentials(),
+        )
+        mock_rest_transport_cls.assert_not_called()
+
+        stub = transport.{{ method.transport_safe_name|snake_case }}
+        result1 = await stub({{ method.input.ident }}())
+        result2 = await stub({{ method.input.ident }}())
+
+        assert result1 is expected_session
+        assert result2 is expected_session
+        mock_rest_transport_cls.assert_called_once_with(
+            host=transport._host,
+            credentials=transport._credentials,
+            client_info=transport._client_info,
+        )
+        assert mock_rest_transport.{{ method.transport_safe_name|snake_case }}.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_{{ service.name|snake_case }}_{{ method.name|snake_case }}_grpc_asyncio_without_async_rest_error():
+    with mock.patch.object(transports.grpc_asyncio, "HAS_ASYNC_REST", False):
+        transport = transports.{{ service.name }}GrpcAsyncIOTransport(
+            credentials=async_anonymous_credentials(),
+        )
+        with pytest.raises(NotImplementedError) as exc_info:
+            await transport.{{ method.transport_safe_name|snake_case }}({{ method.input.ident }}())
+        assert "Async REST transport is required for async resumable upload methods." in str(exc_info.value)
 {% endfor %}
 {% endif %}
 
@@ -1556,13 +1600,40 @@ def test_{{ service.name|snake_case }}_{{ method.name|snake_case }}_rest_start_e
         transport="rest",
     )
     with mock.patch.object(type(client.transport._session), "request") as req:
-        response_value = mock.MagicMock()
-        response_value.status_code = 200
-        response_value.headers = {"x-goog-upload-status": "active", "x-goog-upload-url": "http://localhost/upload"}
-        response_value.content = b""
-        req.return_value = response_value
-        response = client.{{ method.client_method_name|snake_case }}({{ method.input.ident }}())
+        start_response = mock.MagicMock()
+        start_response.ok = True
+        start_response.status_code = 200
+        start_response.headers = {"x-goog-upload-status": "active", "x-goog-upload-url": "http://localhost/upload"}
+        start_response.content = b""
+
+        expected_result = {{ method.output.ident }}(
+            {% for field in method.output.fields.values() | rejectattr('message') %}
+            {% if not field.oneof or field.proto3_optional %}
+            {{ field.name }}={{ field.mock_value }},
+            {% endif %}
+            {% endfor %}
+        )
+        {% if method.output.ident.is_proto_plus_type %}
+        json_return_value = {{ method.output.ident }}.to_json(expected_result)
+        {% else %}
+        json_return_value = json_format.MessageToJson(expected_result)
+        {% endif %}
+        upload_response = mock.MagicMock()
+        upload_response.ok = True
+        upload_response.status_code = 200
+        upload_response.headers = {"x-goog-upload-status": "final"}
+        upload_response.content = json_return_value.encode("UTF-8")
+
+        req.side_effect = [start_response, upload_response]
+        response = client.{{ method.client_method_name|snake_case }}(
+            {{ method.input.ident }}(),
+            config=resumable_transfer.ResumableUploadConfig(headers={"x-custom": "val"}),
+            timeout=10.0,
+        )
         assert isinstance(response, {{ method.client_output.ident }})
+        result = response.upload(b"test data")
+        assert isinstance(result, {{ method.output.ident }})
+        assert result == expected_result
 
 
 {% if rest_async_io_enabled %}
@@ -1574,14 +1645,45 @@ async def test_{{ service.name|snake_case }}_{{ method.name|snake_case }}_rest_a
         credentials=async_anonymous_credentials(),
         transport="rest_asyncio",
     )
-    with mock.patch.object(type(client.transport._session), "request", new_callable=mock.AsyncMock) as req:
-        response_value = mock.MagicMock()
-        response_value.status_code = 200
-        response_value.headers = {"x-goog-upload-status": "active", "x-goog-upload-url": "http://localhost/upload"}
-        response_value.read = mock.AsyncMock(return_value=b"")
-        req.return_value = response_value
-        response = await client.{{ method.client_method_name|snake_case }}({{ method.input.ident }}())
+    with mock.patch.object(type(client.transport._session), "request", new_callable=mock.MagicMock) as req:
+        start_response = mock.MagicMock()
+        start_response.status = 200
+        start_response.headers = {"x-goog-upload-status": "active", "x-goog-upload-url": "http://localhost/upload"}
+        start_response.read = mock.AsyncMock(return_value=b"")
+        start_ctx = mock.MagicMock()
+        start_ctx.__aenter__ = mock.AsyncMock(return_value=start_response)
+        start_ctx.__aexit__ = mock.AsyncMock(return_value=None)
+
+        expected_result = {{ method.output.ident }}(
+            {% for field in method.output.fields.values() | rejectattr('message') %}
+            {% if not field.oneof or field.proto3_optional %}
+            {{ field.name }}={{ field.mock_value }},
+            {% endif %}
+            {% endfor %}
+        )
+        {% if method.output.ident.is_proto_plus_type %}
+        json_return_value = {{ method.output.ident }}.to_json(expected_result)
+        {% else %}
+        json_return_value = json_format.MessageToJson(expected_result)
+        {% endif %}
+        upload_response = mock.MagicMock()
+        upload_response.status = 200
+        upload_response.headers = {"x-goog-upload-status": "final"}
+        upload_response.read = mock.AsyncMock(return_value=json_return_value.encode("UTF-8"))
+        upload_ctx = mock.MagicMock()
+        upload_ctx.__aenter__ = mock.AsyncMock(return_value=upload_response)
+        upload_ctx.__aexit__ = mock.AsyncMock(return_value=None)
+
+        req.side_effect = [start_ctx, upload_ctx]
+        response = await client.{{ method.client_method_name|snake_case }}(
+            {{ method.input.ident }}(),
+            config=resumable_transfer.ResumableUploadConfig(headers={"x-custom": "val"}),
+            timeout=10.0,
+        )
         assert isinstance(response, {{ method.client_output_async.ident }})
+        result = await response.upload(b"test data")
+        assert isinstance(result, {{ method.output.ident }})
+        assert result == expected_result
 {% endif %}
 {% endfor %}
 {% endif %}

packages/gapic-generator/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2

@@ -45,6 +45,8 @@ def test_{{ method_name }}(request_type, transport: str = 'grpc'):
         call.return_value = operations_pb2.Operation(name='operations/spam')
         {% elif method.server_streaming %}
         call.return_value = iter([{{ method.output.ident }}()])
+        {% elif method.is_resumable_upload %}
+        call.return_value = {{ method.client_output.ident }}()
         {% else %}
         call.return_value = {{ method.output.ident }}(
           {% for field in method.output.fields.values() | rejectattr('message')%}
@@ -315,6 +317,10 @@ async def test_{{ method_name }}_async(request_type, transport: str = 'grpc_asyn
         {% elif method.client_streaming and method.server_streaming %}
         call.return_value = mock.Mock(aio.StreamStreamCall, autospec=True)
         call.return_value.read = mock.AsyncMock(side_effect=[{{ method.output.ident }}()])
+        {% elif method.is_resumable_upload %}
+        call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(
+            {{ method.client_output_async.ident }}()
+        )
         {% else %}
         call.return_value ={{ '' }}
             {%- if not method.client_streaming and not method.server_streaming -%}
@@ -1837,20 +1843,37 @@ def test_initialize_client_w_{{transport_name}}():
     request = request_type(**request_init)
 
     # Mock the http request call within the method and fake a BadRequest error.
-    with mock.patch.object({{mocked_session}}, 'request') as req, pytest.raises(core_exceptions.BadRequest):
+    with mock.patch.object({{mocked_session}}, 'request'{% if is_async and method.is_resumable_upload %}, new_callable=mock.MagicMock{% endif %}) as req, pytest.raises(core_exceptions.BadRequest):
         # Wrap the value into a proper Response obj
+        {% if method.is_resumable_upload %}
+        response_value = mock.MagicMock()
+        {% else %}
         response_value = mock.Mock()
+        {% endif %}
         {% if is_async %}
         response_value.read = mock.AsyncMock(return_value=b'{}')
+        {% if method.is_resumable_upload %}
+        response_value.status = 400
+        response_value.__aenter__ = mock.AsyncMock(return_value=response_value)
+        response_value.__aexit__ = mock.AsyncMock(return_value=None)
+        {% endif %}
         {% else %}
         json_return_value = ''
         response_value.json = mock.Mock(return_value={})
+        {% if method.is_resumable_upload %}
+        response_value.ok = False
+        {% endif %}
         {% endif %}{# if is_async #}
         response_value.status_code = 400
         response_value.request = mock.Mock()
         req.return_value = response_value
         req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"}
+        {% if method.is_resumable_upload %}
+        response = {{ await_prefix }}client.{{ method_name }}(request)
+        {{ await_prefix }}response.upload(b"")
+        {% else %}
         {{ await_prefix }}client.{{ method_name }}(request)
+        {% endif %}
 
 {% endif %}{# if 'grpc' in transport #}
 {% endmacro %}
@@ -2272,13 +2295,13 @@ def test_initialize_client_w_{{transport_name}}():
         {% if method.lro %}
         mock.patch.object(operation.Operation, "_set_result_from_operation"), \
         {% endif %}
-        {% if not method.void %}
+        {% if not method.void and not method.is_resumable_upload %}
         mock.patch.object(transports.{{async_method_prefix}}{{ service.name }}RestInterceptor, "post_{{method.name|snake_case}}") as post, \
         mock.patch.object(transports.{{async_method_prefix}}{{ service.name }}RestInterceptor, "post_{{method.name|snake_case}}_with_metadata") as post_with_metadata, \
         {% endif %}
         mock.patch.object(transports.{{async_method_prefix}}{{ service.name }}RestInterceptor, "pre_{{ method.name|snake_case }}") as pre:
         pre.assert_not_called()
-        {% if not method.void %}
+        {% if not method.void and not method.is_resumable_upload %}
         post.assert_not_called()
         post_with_metadata.assert_not_called()
         {% endif %}
@@ -2297,7 +2320,7 @@ def test_initialize_client_w_{{transport_name}}():
         req.return_value = mock.Mock()
         req.return_value.status_code = 200
         req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"}
-        {% if not method.void %}
+        {% if not method.void and not method.is_resumable_upload %}
         return_value = {% if method.output.ident.is_proto_plus_type %}{{ method.output.ident }}.to_json({{ method.output.ident }}()){% else %}json_format.MessageToJson({{ method.output.ident }}()){% endif %}
 
         {% if method.server_streaming %}
@@ -2313,23 +2336,23 @@ def test_initialize_client_w_{{transport_name}}():
         req.return_value.content = return_value
         {% endif %}{# end is_async #}
         {% endif %}{# end method.server_streaming #}
-        {% endif %}{# end not method.void #}
+        {% endif %}{# end not method.void and not method.is_resumable_upload #}
 
         request = {{ method.input.ident }}()
         metadata =[
             ("key", "val"),
             ("cephalopod", "squid"),
         ]
         pre.return_value = request, metadata
-        {% if not method.void %}
+        {% if not method.void and not method.is_resumable_upload %}
         post.return_value = {{ method.output.ident }}()
         post_with_metadata.return_value = {{ method.output.ident }}(), metadata
         {% endif %}
 
         {{await_prefix}}client.{{ method_name }}(request, metadata=[("key", "val"), ("cephalopod", "squid"),])
 
         pre.assert_called_once()
-        {% if not method.void %}
+        {% if not method.void and not method.is_resumable_upload %}
         post.assert_called_once()
         post_with_metadata.assert_called_once()
         {% endif %}