Repository: jina-ai/serve
Pull request: refactor: merge grpc and rest client
This PR aims at simplify C/S mode (Flow as service) Jina Python client usage with a uniform facade, before and after:

```python
# before
from jina import Client
from jina.clients import WebsocketClient

client = Client('localhost', 12345) # connect to grpc gateway
client = WebsocketClient('localhost', 12345) # connect to restful gateway

# now
from jina import Client
client = Client('localhost', 12345)  # connect to grpc gateway
client = Client('localhost', 12345, restful=True)  # connect to restful gateway

# it uses the same semantics as what we have defined for the flow:
f = Flow()
f = Flow(restful=True)
```

Problem statement:
Merge jina restful and grpc client usage 
```python
# before
from jina import Client
from jina.clients import WebsocketClient

client = Client('localhost', 12345) # connect to grpc gateway
client = WebsocketClient('localhost', 12345) # connect to restful gateway

# expected
from jina import Client
client = Client('localhost', 12345)  # connect to grpc gateway
client = Client('localhost', 12345, restful=True)  # connect to restful gateway
```

Hints from issue discussion:
cc @CatStark

Commit message: feat: add grpc client back to distributed test

Patch to review:
diff --git a/.github/2.0/cookbooks/Flow.md b/.github/2.0/cookbooks/Flow.md
index 818cd1b213fad..60c03f9262287 100644
--- a/.github/2.0/cookbooks/Flow.md
+++ b/.github/2.0/cookbooks/Flow.md
@@ -717,7 +717,7 @@ c.post('/')
 ```
 
 ```text
-         Client@27219[S]:connected to the gateway at 192.168.1.14:12345!
+         GrpcClient@27219[S]:connected to the gateway at 192.168.1.14:12345!
   |█                   | 📃    100 ⏱️ 0.0s 🐎 26690.1/s      1   requests takes 0 seconds (0.00s)
 	✅ done in ⏱ 0 seconds 🐎 24854.8/s
 ```
@@ -834,17 +834,17 @@ When use `curl`, make sure to pass the `-N/--no-buffer` flag.
 
 ### Python Client with REST Request
 
-In some case when `restful=True`, you may still need a Python client to query the server for debugging. You can use `WebsocketClient` in this case.
+In some case when `restful=True`, you may still need a Python client to query the server for debugging. You can use `Client` with an additional parameter `restful` in this case.
 
 ```python
-from jina.clients import WebSocketClient
+from jina import Client
 
-c = WebSocketClient(host='192.168.1.14', port_expose=12345)
+c = Client(host='192.168.1.14', port_expose=12345, restful=True)
 c.post('/')
 ```
 
 ```text
-WebSocketClient@27622[S]:Connected to the gateway at 192.168.1.14:12345
+         WebSocketClient@27622[S]:Connected to the gateway at 192.168.1.14:12345
   |█                   | 📃    100 ⏱️ 0.0s 🐎 19476.6/s      1   requests takes 0 seconds (0.00s)
 	✅ done in ⏱ 0 seconds 🐎 18578.9/s
 ```
diff --git a/jina/clients/__init__.py b/jina/clients/__init__.py
index abe4dceaad6fd..6a154f3c66f2e 100644
--- a/jina/clients/__init__.py
+++ b/jina/clients/__init__.py
@@ -4,25 +4,42 @@
 from .mixin import PostMixin
 from .request import GeneratorSourceType
 from .websocket import WebSocketClientMixin
+from ..parsers import set_client_cli_parser
 
-__all__ = ['Client', 'WebSocketClient']
+__all__ = ['Client', 'GrpcClient', 'WebSocketClient']
 
 
-class Client(PostMixin, BaseClient):
+def Client(host: str, port_expose: int, restful: bool = False) -> 'BaseClient':
+    """Jina Python client.
+
+    :param host: Host address of the flow.
+    :param port_expose: Port number of the flow.
+    :param restful: If connect to a Restful gateway, default is ``False``, connect to GrpcGateway.
+    :return: An instance of :class:`GrpcClient` or :class:`WebSocketClient`.
+    """
+    args = set_client_cli_parser().parse_args(
+        ['--host', host, '--port-expose', str(port_expose)]
+    )
+    if restful:
+        return WebSocketClient(args)
+    else:
+        return GrpcClient(args)
+
+
+class GrpcClient(PostMixin, BaseClient):
     """A simple Python client for connecting to the gRPC gateway.
 
     It manages the asyncio event loop internally, so all interfaces are synchronous from the outside.
     """
 
     @property
-    def client(self) -> 'Client':
+    def client(self) -> 'GrpcClient':
         """Return the client object itself
-
         .. # noqa: DAR201"""
         return self
 
 
-class WebSocketClient(Client, WebSocketClientMixin):
+class WebSocketClient(GrpcClient, WebSocketClientMixin):
     """A Python Client to stream requests from a Flow with a REST Gateway.
 
     :class:`WebSocketClient` shares the same interface as :class:`Client` and provides methods like
diff --git a/jina/clients/mixin.py b/jina/clients/mixin.py
index 171f21f13502e..8b2d1c9fa3827 100644
--- a/jina/clients/mixin.py
+++ b/jina/clients/mixin.py
@@ -7,7 +7,7 @@
 
 
 class PostMixin:
-    """The Post Mixin class for Client and Flow """
+    """The Post Mixin class for Client and Flow"""
 
     def post(
         self,
@@ -63,7 +63,7 @@ async def _get_results(*args, **kwargs):
 
 
 class AsyncPostMixin:
-    """The Async Post Mixin class for AsyncClient and AsyncFlow """
+    """The Async Post Mixin class for AsyncClient and AsyncFlow"""
 
     async def post(
         self,
diff --git a/jina/flow/base.py b/jina/flow/base.py
index f479dd6102062..d20a21ed47418 100644
--- a/jina/flow/base.py
+++ b/jina/flow/base.py
@@ -12,7 +12,8 @@
 
 from .builder import build_required, _build_flow, _hanging_pods
 from .. import __default_host__
-from ..clients import Client, WebSocketClient
+from ..clients.base import BaseClient
+from ..clients import GrpcClient, WebSocketClient
 from ..enums import FlowBuildLevel, PodRoleType, FlowInspectType
 from ..excepts import FlowTopologyError, FlowMissingPodError
 from ..helper import (
@@ -68,7 +69,9 @@ class BaseFlow(JAMLCompatible, ExitStack, metaclass=FlowType):
     :param env: environment variables shared by all Pods
     """
 
-    _cls_client = Client  #: the type of the Client, can be changed to other class
+    _cls_client = (
+        GrpcClient  #: the type of the GrpcClient, can be changed to other class
+    )
 
     def __init__(
         self,
@@ -169,7 +172,7 @@ def _add_gateway(self, needs, **kwargs):
                 name=pod_name,
                 ctrl_with_ipc=True,  # otherwise ctrl port would be conflicted
                 runtime_cls='GRPCRuntime'
-                if self._cls_client == Client
+                if self._cls_client == GrpcClient
                 else 'RESTRuntime',
                 pod_role=PodRoleType.GATEWAY,
                 identity=self.args.identity,
@@ -601,8 +604,8 @@ def __eq__(self, other: 'BaseFlow') -> bool:
 
     @property
     @build_required(FlowBuildLevel.GRAPH)
-    def client(self) -> 'Client':
-        """Return a :class:`Client` object attach to this Flow.
+    def client(self) -> 'BaseClient':
+        """Return a :class:`BaseClient` object attach to this Flow.
 
         .. # noqa: DAR201"""
         kwargs = {}
@@ -617,7 +620,6 @@ def client(self) -> 'Client':
         # show progress when client is used inside the flow, for better log readability
         if 'show_progress' not in kwargs:
             args.show_progress = True
-
         return self._cls_client(args)
 
     @property
@@ -894,7 +896,7 @@ def _show_success_message(self):
         self.logger.info('\n' + '\n'.join(address_table))
 
     def block(self):
-        """Block the process until user hits KeyboardInterrupt """
+        """Block the process until user hits KeyboardInterrupt"""
         try:
             threading.Event().wait()
         except KeyboardInterrupt:
@@ -912,7 +914,7 @@ def use_grpc_gateway(self, port: Optional[int] = None):
 
     def _switch_gateway(self, gateway: str, port: int):
         restful = gateway == 'RESTRuntime'
-        client = WebSocketClient if gateway == 'RESTRuntime' else Client
+        client = WebSocketClient if gateway == 'RESTRuntime' else GrpcClient
 
         # globally register this at Flow level
         self._cls_client = client
diff --git a/tests/distributed/test_join_local_from_remote/test_integration.py b/tests/distributed/test_join_local_from_remote/test_integration.py
index 672bf96aa0dee..e5e86aacffb9a 100644
--- a/tests/distributed/test_join_local_from_remote/test_integration.py
+++ b/tests/distributed/test_join_local_from_remote/test_integration.py
@@ -3,7 +3,7 @@
 import pytest
 
 from jina import Document
-from jina.clients import Client
+from jina.clients import Client, GrpcClient
 from jina.parsers import set_client_cli_parser
 from tests import validate_callback
 from ..helpers import create_flow, assert_request
@@ -23,16 +23,26 @@ def doc_to_index():
 
 @pytest.fixture
 def client():
+    return Client(host='localhost', port_expose=45678)
+
+
+@pytest.fixture
+def grpc_client():
     args = set_client_cli_parser().parse_args(
         ['--host', 'localhost', '--port-expose', '45678']
     )
 
-    return Client(args)
+    return GrpcClient(args)
+
+
+@pytest.fixture(params=['client', 'grpc_client'])
+def client_instance(request):
+    return request.getfixturevalue(request.param)
 
 
 @pytest.mark.timeout(360)
 @pytest.mark.parametrize('docker_compose', [compose_yml], indirect=['docker_compose'])
-def test_flow(docker_compose, doc_to_index, client, mocker):
+def test_flow(docker_compose, doc_to_index, client_instance, mocker):
     def validate_resp(resp):
         assert len(resp.data.docs) == 2
         assert resp.data.docs[0].text == 'test'
@@ -41,7 +51,7 @@ def validate_resp(resp):
     mock = mocker.Mock()
     flow_id = create_flow(flow_yaml=flow_yaml, pod_dir=os.path.join(cur_dir, 'pods'))
 
-    client.search(inputs=[doc_to_index], on_done=mock)
+    client_instance.search(inputs=[doc_to_index], on_done=mock)
 
     assert_request(method='get', url=f'http://localhost:8000/flows/{flow_id}')
 
diff --git a/tests/unit/clients/python/test_client.py b/tests/unit/clients/python/test_client.py
index f242fd4671904..5a8614e584d2b 100644
--- a/tests/unit/clients/python/test_client.py
+++ b/tests/unit/clients/python/test_client.py
@@ -6,7 +6,7 @@
 
 from jina import Executor, DocumentArray, requests as req
 from jina import helper, Document
-from jina.clients import Client, WebSocketClient
+from jina.clients import Client
 from jina.excepts import BadClientInput
 from jina.flow import Flow
 from jina.parsers import set_gateway_parser, set_client_cli_parser
@@ -37,19 +37,26 @@ def test_img_2():
     return 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAICAIAAABLbSncAAAA2ElEQVR4nADIADf/AvdGjTZeOlQq07xSYPgJjlWRwfWEBx2+CgAVrPrP+O5ghhOa+a0cocoWnaMJFAsBuCQCgiJOKDBcIQTiLieOrPD/cp/6iZ/Iu4HqAh5dGzggIQVJI3WqTxwVTDjs5XJOy38AlgHoaKgY+xJEXeFTyR7FOfF7JNWjs3b8evQE6B2dTDvQZx3n3Rz6rgOtVlaZRLvR9geCAxuY3G+0mepEAhrTISES3bwPWYYi48OUrQOc//IaJeij9xZGGmDIG9kc73fNI7eA8VMBAAD//0SxXMMT90UdAAAAAElFTkSuQmCC'
 
 
+@pytest.fixture(scope='function')
+def args():
+    return set_client_cli_parser()
+
+
 @pytest.mark.parametrize(
     'inputs', [iter([b'1234', b'45467']), iter([DocumentProto(), DocumentProto()])]
 )
-def test_check_input_success(inputs):
-    Client.check_input(inputs)
+def test_check_input_success(inputs, args):
+    client = Client(host='localhost', port_expose=12345)
+    client.check_input(inputs)
 
 
 @pytest.mark.parametrize(
     'inputs', [iter([list(), list(), [12, 2, 3]]), iter([set(), set()])]
 )
-def test_check_input_fail(inputs):
+def test_check_input_fail(inputs, args):
+    client = Client(host='localhost', port_expose=12345)
     with pytest.raises(BadClientInput):
-        Client.check_input(inputs)
+        client.check_input(inputs)
 
 
 @pytest.mark.parametrize(
@@ -122,18 +129,13 @@ def test_client_csv(restful, mocker, func_name):
 
 # Timeout is necessary to fail in case of hanging client requests
 @pytest.mark.timeout(5)
-def test_client_websocket(mocker, flow_with_rest_api_enabled):
+def test_client_websocket(mocker, flow_with_rest_api_enabled, args):
     with flow_with_rest_api_enabled:
         time.sleep(0.5)
-        client = WebSocketClient(
-            set_client_cli_parser().parse_args(
-                [
-                    '--host',
-                    'localhost',
-                    '--port-expose',
-                    str(flow_with_rest_api_enabled.port_expose),
-                ]
-            )
+        client = Client(
+            host='localhost',
+            port_expose=str(flow_with_rest_api_enabled.port_expose),
+            restful=True,
         )
         # Test that a regular index request triggers the correct callbacks
         on_always_mock = mocker.Mock()
@@ -158,16 +160,16 @@ def test_client_websocket(mocker, flow_with_rest_api_enabled):
         mock.assert_not_called()
 
 
-@pytest.mark.parametrize('client_cls', [Client, WebSocketClient])
-def test_client_from_kwargs(client_cls):
-    client_cls(port_expose=12345, host='0.0.0.1')
+def test_client_from_kwargs():
+    Client(port_expose=12345, host='0.0.0.1')
+    Client(port_expose=12345, host='0.0.0.1', restful=True)
 
 
 def test_independent_client():
     with Flow() as f:
-        c = Client(port_expose=f.port_expose)
+        c = Client(host='localhost', port_expose=f.port_expose)
         c.post('/')
 
     with Flow(restful=True) as f:
-        c = WebSocketClient(port_expose=f.port_expose)
+        c = Client(host='localhost', port_expose=f.port_expose, restful=True)
         c.post('/')
diff --git a/tests/unit/clients/python/test_io.py b/tests/unit/clients/python/test_io.py
index 5254f4c62a617..1bb1c68da662d 100644
--- a/tests/unit/clients/python/test_io.py
+++ b/tests/unit/clients/python/test_io.py
@@ -16,6 +16,11 @@
 cur_dir = os.path.dirname(os.path.abspath(__file__))
 
 
+@pytest.fixture(scope='function')
+def client():
+    return Client(host='localhost', port_expose=123456)
+
+
 @pytest.fixture(scope='function')
 def filepath(tmpdir):
     input_filepath = os.path.join(tmpdir, 'input_file.csv')
@@ -148,8 +153,8 @@ def test_input_lines_with_jsonlines_docs_groundtruth():
         ('*.*', True, None, 0.5, None),
     ],
 )
-def test_input_files(patterns, recursive, size, sampling_rate, read_mode):
-    Client.check_input(
+def test_input_files(patterns, recursive, size, sampling_rate, read_mode, client):
+    client.check_input(
         from_files(
             patterns=patterns,
             recursive=recursive,
@@ -160,13 +165,13 @@ def test_input_files(patterns, recursive, size, sampling_rate, read_mode):
     )
 
 
-def test_input_files_with_invalid_read_mode():
+def test_input_files_with_invalid_read_mode(client):
     with pytest.raises(BadClientInput):
-        Client.check_input(from_files(patterns='*.*', read_mode='invalid'))
+        client.check_input(from_files(patterns='*.*', read_mode='invalid'))
 
 
 @pytest.mark.parametrize(
     'array', [np.random.random([100, 4, 2]), ['asda', 'dsadas asdasd']]
 )
-def test_input_numpy(array):
-    Client.check_input(from_ndarray(array))
+def test_input_numpy(array, client):
+    client.check_input(from_ndarray(array))


Code review comments:
diff_hunk: @@ -912,7 +914,7 @@ def use_grpc_gateway(self, port: Optional[int] = None):
 
     def _switch_gateway(self, gateway: str, port: int):
         restful = gateway == 'RESTRuntime'
-        client = WebSocketClient if gateway == 'RESTRuntime' else Client
+        client = WebSocketClient if gateway == 'RESTRuntime' else GrpcClient
path: jina/flow/base.py
text: @user1:
Why can we not use `Client` here passing restful? 

@author:
@user1 we can, while it does not make a lot difference since `Client` is actually initialising `GrpcClient` and `WebSocketClient`

@user2:
then lets have a single poijt of entry, change it to Client with the proper argument, u can use partial functions if needed

@user1:
I meant, not using `GrpcClient` & `WebSocketClient` at all in `flow/base.py` and using `Client` directly. Not important though.

@author:
gave a try, and I broke a lot of tests. I'll look into if I can optimise this part in a future PR.

@user3:
> and I broke a lot of tests

I believe the reason was because your `Client()` function does not accept the full arguments that CLI defined (as I mentioned here https://github.com/jina-ai/jina/pull/2565#discussion_r647896285). This oversimplification causes incomplete setup of the client and yields the broken test.