Repository: jina-ai/serve
Pull request: test: compound pod sockets
testing the socket configuration of the compound pod
closes https://github.com/jina-ai/jina/issues/2348

Problem statement:
Add unittests of CompoundPod showing ports/hosts and sockets configuration
**Describe the feature**
Add a unittest in `CompoundPod` to test heavily which ports/hosts/socket configurations are set.

Right now, The argument parsing of these components is not so clean, and having such a test, would help refactoring in a safer way in the future

Test socket configuration of the compound pod
The compound pod will be refactored. To make sure the refactoring does not introduce bugs, we keep on adding tests.

Hints from issue discussion:

Commit message: test: compound pod sockets

Patch to review:
diff --git a/tests/unit/peapods/pods/test_compound_pods.py b/tests/unit/peapods/pods/test_compound_pods.py
index a05f54ace5852..f56207e6dcbf6 100644
--- a/tests/unit/peapods/pods/test_compound_pods.py
+++ b/tests/unit/peapods/pods/test_compound_pods.py
@@ -1,7 +1,8 @@
 import pytest
 
+from jina.enums import SchedulerType, SocketType, PollingType
 from jina.parsers import set_pod_parser
-from jina.peapods import CompoundPod
+from jina.peapods import CompoundPod, Pod
 
 
 @pytest.fixture(scope='function')
@@ -162,3 +163,65 @@ def test_pod_naming_with_parallel(runtime):
         assert bp.replica_list[2].peas[1].runtime.name == 'pod/2/tail/ZEDRuntime'
         assert bp.replica_list[2].peas[2].runtime.name == 'pod/2/0/ZEDRuntime'
         assert bp.replica_list[2].peas[3].runtime.name == 'pod/2/1/ZEDRuntime'
+
+
+@pytest.mark.parametrize(
+    'polling, parallel, pea_scheduling, pea_socket_in, pea_socket_out',
+    (
+        (
+            'all',
+            1,
+            SchedulerType.LOAD_BALANCE,
+            SocketType.SUB_CONNECT,
+            SocketType.PUSH_CONNECT,
+        ),
+        (
+            'all',
+            2,
+            SchedulerType.LOAD_BALANCE,
+            SocketType.SUB_CONNECT,
+            SocketType.PUSH_CONNECT,
+        ),
+        (
+            'any',
+            2,
+            SchedulerType.LOAD_BALANCE,
+            SocketType.DEALER_CONNECT,
+            SocketType.PUSH_CONNECT,
+        ),
+    ),
+)
+def test_sockets(polling, parallel, pea_scheduling, pea_socket_in, pea_socket_out):
+    polling_type = PollingType.ALL if polling == 'all' else PollingType.ANY
+    args = set_pod_parser().parse_args(
+        [
+            '--name',
+            'pod',
+            '--parallel',
+            f'{parallel}',
+            '--polling',
+            f'{polling}',
+            '--replicas',
+            '3',
+        ]
+    )
+    compound_pod = CompoundPod(args)
+    replica_args = compound_pod.replicas_args
+    head = replica_args['head']
+    assert head.socket_in == SocketType.PULL_BIND
+    assert head.socket_out == SocketType.ROUTER_BIND
+    assert head.scheduling == SchedulerType.LOAD_BALANCE
+    tail = replica_args['tail']
+    assert tail.socket_in == SocketType.PULL_BIND
+    assert tail.socket_out == SocketType.PUSH_BIND
+    assert tail.scheduling == SchedulerType.LOAD_BALANCE
+    replicas = replica_args['replicas']
+    for pod_args in replicas:
+        pod = Pod(pod_args)
+        if parallel > 1:
+            assert pod_args.polling == polling_type
+            for pea in pod.peas_args['peas']:
+                assert pea.polling == polling_type
+                assert pea.socket_in == pea_socket_in
+                assert pea.socket_out == pea_socket_out
+                assert pea.scheduling == pea_scheduling


Code review comments:
diff_hunk: @@ -162,3 +163,65 @@ def test_pod_naming_with_parallel(runtime):
         assert bp.replica_list[2].peas[1].runtime.name == 'pod/2/tail/ZEDRuntime'
         assert bp.replica_list[2].peas[2].runtime.name == 'pod/2/0/ZEDRuntime'
         assert bp.replica_list[2].peas[3].runtime.name == 'pod/2/1/ZEDRuntime'
+
+
+@pytest.mark.parametrize(
+    'polling, parallel, pea_scheduling, pea_socket_in, pea_socket_out',
+    (
+        (
+            'all',
+            1,
+            SchedulerType.LOAD_BALANCE,
+            SocketType.SUB_CONNECT,
+            SocketType.PUSH_CONNECT,
+        ),
+        (
+            'all',
+            2,
+            SchedulerType.LOAD_BALANCE,
+            SocketType.SUB_CONNECT,
+            SocketType.PUSH_CONNECT,
+        ),
+        (
+            'any',
+            2,
+            SchedulerType.LOAD_BALANCE,
+            SocketType.DEALER_CONNECT,
+            SocketType.PUSH_CONNECT,
+        ),
+    ),
+)
+def test_sockets(polling, parallel, pea_scheduling, pea_socket_in, pea_socket_out):
+    polling_type = PollingType.ALL if polling == 'all' else PollingType.ANY
+    args = set_pod_parser().parse_args(
+        [
+            '--name',
+            'pod',
+            '--parallel',
+            f'{parallel}',
+            '--polling',
+            f'{polling}',
+            '--replicas',
+            '3',
+        ]
+    )
+    compound_pod = CompoundPod(args)
+    replica_args = compound_pod.replicas_args
+    head = replica_args['head']
+    assert head.socket_in == SocketType.PULL_BIND
+    assert head.socket_out == SocketType.ROUTER_BIND
+    assert head.scheduling == SchedulerType.LOAD_BALANCE
+    tail = replica_args['tail']
+    assert tail.socket_in == SocketType.PULL_BIND
+    assert tail.socket_out == SocketType.PUSH_BIND
+    assert tail.scheduling == SchedulerType.LOAD_BALANCE
+    replicas = replica_args['replicas']
+    for pod_args in replicas:
+        pod = Pod(pod_args)
path: tests/unit/peapods/pods/test_compound_pods.py
text: @user1:
I see why this is needed, but this is a code smell. We should not need to add logic in a test. We should be able to instantiate a CompoundPod and test what's inside their Pods and Peas.

I think what the test is showing is the need to change the way we initialize our layers.

We should have at `__init__` time:

```
CompoundPod instantiates inner Pods
Pod instantiates inner Peas
Peas instantiate inner Runtime (Already so)
```

At `__enter__` time:
```
CompoundPod enters the context of Pod
Pod enters the context of Pea
Pea starts runtime
```



@author:
I see your point. There should be a differentiation between the configuration and starting of compound pods, pods and peas.
The refactoring is not trivial, though.
It does involve refactoring the `noblock_on_start` attribute used by compound pod, pod, and pea.
I would suggest doing this in a different pr

@user1:
The noblock on start can be set at the Init also, it only affects the initialization right? or is a new attribute to keep.

if u want to do in another PR I would suggest to do before this test. 

Because this test has too much logic in it

@user1:
PR with proposal in #2360