chore: fallback to unary_unary when wrapping async callables (#653)

daniel-sanche · web-flow · commit c9e1cbb3297f · 2024-06-04T07:40:25.000-04:00
* fix: fallback to unary_unary when wrapping async callables We recently [made a change](ab22afd) to move some computation in the async rpc wrapper from call-time to wrap-time. This way, individual calls would execute faster, since they don't have to re-compute some data on each call A side-effect of this change is that now some [type validation](https://github.com/googleapis/python-api-core/blob/d96eb5cdd8120bfec97d62b09512c6fecc325be8/google/api_core/grpc_helpers_async.py#L209) happens earlier. This caused some downstream tests to fail when a mock grpc channel is used. The wrapper doesn't know how to handle the mock.Mock type, and raises an exception while constructing the client object This PR fixes the issue by falling back to the unary wrapper when the callable type is unknown, rather than raising an exception. This is in-line with how [the sync version handles it](https://github.com/googleapis/python-api-core/blob/d96eb5cdd8120bfec97d62b09512c6fecc325be8/google/api_core/grpc_helpers.py#L198) * fixed elif to if * removed outdated test
diff --git a/google/api_core/grpc_helpers_async.py b/google/api_core/grpc_helpers_async.py
@@ -197,16 +197,15 @@ def wrap_errors(callable_):
     Returns: Callable: The wrapped gRPC callable.
     """
     grpc_helpers._patch_callable_name(callable_)
-    if isinstance(callable_, aio.UnaryUnaryMultiCallable):
-        return _wrap_unary_errors(callable_)
-    elif isinstance(callable_, aio.UnaryStreamMultiCallable):
+
+    if isinstance(callable_, aio.UnaryStreamMultiCallable):
         return _wrap_stream_errors(callable_, _WrappedUnaryStreamCall)
     elif isinstance(callable_, aio.StreamUnaryMultiCallable):
         return _wrap_stream_errors(callable_, _WrappedStreamUnaryCall)
     elif isinstance(callable_, aio.StreamStreamMultiCallable):
         return _wrap_stream_errors(callable_, _WrappedStreamStreamCall)
     else:
-        raise TypeError("Unexpected type of callable: {}".format(type(callable_)))
+        return _wrap_unary_errors(callable_)
 
 
 def create_channel(
diff --git a/tests/asyncio/test_grpc_helpers_async.py b/tests/asyncio/test_grpc_helpers_async.py
@@ -165,19 +165,6 @@ async def test_wrap_stream_errors_stream_stream():
     assert mock_call.wait_for_connection.call_count == 1
 
 
-@pytest.mark.asyncio
-async def test_wrap_errors_type_error():
-    """
-    If wrap_errors is called with an unexpected type, it should raise a TypeError.
-    """
-    mock_call = mock.Mock()
-    multicallable = mock.Mock(return_value=mock_call)
-
-    with pytest.raises(TypeError) as exc:
-        grpc_helpers_async.wrap_errors(multicallable)
-    assert "Unexpected type" in str(exc.value)
-
-
 @pytest.mark.asyncio
 async def test_wrap_stream_errors_raised():
     grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT)
