Merge pull request #55 from Yoctol/cph-patch-fix-channel-option

Fix channel option not using bug
Yoctol · Feb 21, 2020 · 857ecfb · 857ecfb
2 parents 25d2cd8 + 40cb9c8
commit 857ecfb
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 11 deletions.
diff --git a/serving_utils/client.py b/serving_utils/client.py
@@ -57,10 +57,14 @@ def __init__(
         if channel_options is None:
             channel_options = {}
         if pem is None:
-            make_sync_channel = grpc.insecure_channel
+            make_sync_channel = partial(grpc.insecure_channel, options=channel_options)
         else:
             creds = grpc.ssl_channel_credentials(pem)
-            make_sync_channel = partial(grpc.secure_channel, credentials=creds)
+            make_sync_channel = partial(
+                grpc.secure_channel,
+                credentials=creds,
+                options=channel_options,
+            )
 
         if loop is None:
             loop = asyncio.get_event_loop()
@@ -120,12 +124,6 @@ def __init__(
         if channel_options is None:
             channel_options = {}
         self._channel_options = channel_options
-        if pem is None:
-            make_sync_channel = grpc.insecure_channel
-        else:
-            creds = grpc.ssl_channel_credentials(pem)
-            make_sync_channel = partial(grpc.secure_channel, credentials=creds)
-        self._make_sync_channel = make_sync_channel
 
         if loop is None:
             loop = asyncio.get_event_loop()

diff --git a/serving_utils/tests/test_client.py b/serving_utils/tests/test_client.py
@@ -8,7 +8,7 @@
 import grpc
 import grpc._channel
 import grpclib
-from ..client import Client, RetryFailed
+from ..client import Client, RetryFailed, Connection
 
 import numpy as np
 from serving_utils import PredictInput
@@ -117,8 +117,10 @@ def hostname_resolution_change(mock_gethostbyname_ex, new_addr_list):
           side_effect=create_a_fake_grpclib_channel).start()
     patch('serving_utils.client.grpc.secure_channel',
           side_effect=create_a_fake_grpc_channel).start()
-    patch('serving_utils.client.grpc.insecure_channel',
-          side_effect=create_a_fake_grpc_channel).start()
+    f.patch_insecure_channel = patch(
+        'serving_utils.client.grpc.insecure_channel',
+        side_effect=create_a_fake_grpc_channel,
+    ).start()
     patch('serving_utils.client.prediction_service_grpc.PredictionServiceStub',
           side_effect=create_a_fake_async_stub).start()
     patch('serving_utils.client.prediction_service_pb2_grpc.PredictionServiceStub',
@@ -257,6 +259,16 @@ def test_RetryFailed():
     assert "UNKNOWN" in str(retry_error)
 
 
+def test_Connection_with_channel_options():
+    t = test_Connection_with_channel_options
+    channel_options = [('max_receive_body', 128 * 1000 * 1000)]
+    Connection('127.0.0.1', 8500, channel_options=channel_options)
+    t.patch_insecure_channel.assert_called_once_with(
+        '127.0.0.1:8500',
+        options=channel_options,
+    )
+
+
 @pytest.mark.asyncio
 async def test_model_not_found_error_passes_through_sync_predict():