address review feedback

Author: parthea

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

@@ -381,13 +381,17 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
                         )
                 self._stubs['{{ method.transport_safe_name|snake_case }}'] = _ErrorStub()
             else:
-                if self._rest_transport is None:
-                    self._rest_transport = {{ service.name }}RestTransport(
-                        host=self._host,
-                        credentials=self._credentials,
-                        client_info=self._client_info,
-                    )
-                self._stubs['{{ method.transport_safe_name|snake_case }}'] = self._rest_transport.{{ method.transport_safe_name|snake_case }}
+                transport = self
+                class _RestStub:
+                    def __call__(self, *args, **kwargs):
+                        if transport._rest_transport is None:
+                            transport._rest_transport = {{ service.name }}RestTransport(
+                                host=transport._host,
+                                credentials=transport._credentials,
+                                client_info=transport._client_info,
+                            )
+                        return transport._rest_transport.{{ method.transport_safe_name|snake_case }}(*args, **kwargs)
+                self._stubs['{{ method.transport_safe_name|snake_case }}'] = _RestStub()
             {% else %}
             self._stubs['{{ method.transport_safe_name|snake_case }}'] = self._logged_channel.{{ method.grpc_stub_type }}(
                 '/{{ '.'.join(method.meta.address.package) }}.{{ service.name }}/{{ method.name }}',

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

@@ -1540,11 +1540,19 @@ def test_{{ service.name|snake_case }}_{{ method.name|snake_case }}_grpc_rest_st
     transport = transports.{{ service.name }}GrpcTransport(
         credentials=ga_credentials.AnonymousCredentials(),
     )
+    assert transport._rest_transport is None
+    stub = transport.{{ method.transport_safe_name|snake_case }}
+    assert transport._rest_transport is None
+
+    result1 = stub({{ method.input.ident }}())
     rest_transport = transport._rest_transport
     assert rest_transport is not None
-    transport._stubs.pop("{{ method.transport_safe_name|snake_case }}", None)
-    _ = transport.{{ method.transport_safe_name|snake_case }}
+    assert isinstance(result1, {{ method.client_output.ident }})
+
+    result2 = stub({{ method.input.ident }}())
+    assert isinstance(result2, {{ method.client_output.ident }})
     assert transport._rest_transport is rest_transport
+
     with mock.patch.object(rest_transport, "close") as mock_rest_close:
         transport.close()
         mock_rest_close.assert_called_once()

packages/gapic-generator/tests/integration/goldens/showcase/google/showcase_v1beta1/services/resumable_upload_service/transports/grpc.py

@@ -354,13 +354,17 @@ def __call__(self, *args, **kwargs):
                         )
                 self._stubs['upload_media'] = _ErrorStub()
             else:
-                if self._rest_transport is None:
-                    self._rest_transport = ResumableUploadServiceRestTransport(
-                        host=self._host,
-                        credentials=self._credentials,
-                        client_info=self._client_info,
-                    )
-                self._stubs['upload_media'] = self._rest_transport.upload_media
+                transport = self
+                class _RestStub:
+                    def __call__(self, *args, **kwargs):
+                        if transport._rest_transport is None:
+                            transport._rest_transport = ResumableUploadServiceRestTransport(
+                                host=transport._host,
+                                credentials=transport._credentials,
+                                client_info=transport._client_info,
+                            )
+                        return transport._rest_transport.upload_media(*args, **kwargs)
+                self._stubs['upload_media'] = _RestStub()
         return self._stubs['upload_media']
 
     def close(self):

packages/gapic-generator/tests/integration/goldens/showcase/tests/unit/gapic/showcase_v1beta1/test_resumable_upload_service.py

@@ -2294,11 +2294,19 @@ def test_resumable_upload_service_upload_media_grpc_rest_stub():
     transport = transports.ResumableUploadServiceGrpcTransport(
         credentials=ga_credentials.AnonymousCredentials(),
     )
+    assert transport._rest_transport is None
+    stub = transport.upload_media
+    assert transport._rest_transport is None
+
+    result1 = stub(resumable_upload.UploadMediaRequest())
     rest_transport = transport._rest_transport
     assert rest_transport is not None
-    transport._stubs.pop("upload_media", None)
-    _ = transport.upload_media
+    assert isinstance(result1, resumable_transfer.ResumableUploadSession)
+
+    result2 = stub(resumable_upload.UploadMediaRequest())
+    assert isinstance(result2, resumable_transfer.ResumableUploadSession)
     assert transport._rest_transport is rest_transport
+
     with mock.patch.object(rest_transport, "close") as mock_rest_close:
         transport.close()
         mock_rest_close.assert_called_once()