add unit test

Signed-off-by: Sidhant Kohli <sidhant.kohli@gmail.com>

Author: kohlisid

pynumaflow/mapper/async_multiproc_server.py

@@ -82,7 +82,7 @@ def start(self):
                     bind_address = f"0.0.0.0:{reserved_port}"
                     ports.append(f"http://{bind_address}")
             else:
-                bind_address = f"unix://{self.sock_path}{idx}.sock"
+                bind_address = f"{self.sock_path}{idx}.sock"
             _LOGGER.info("Binding server to: %s", bind_address)
 
             worker = multiprocessing.Process(

tests/map/test_async_multiproc.py

@@ -0,0 +1,83 @@
+import os
+import time
+import uuid
+import unittest
+import grpc
+from multiprocessing import Process
+from pynumaflow.mapper import Datum, Messages, Message
+from pynumaflow.proto.mapper import map_pb2_grpc
+from tests.map.test_async_mapper import request_generator
+from tests.map.utils import get_test_datums
+from pynumaflow.mapper import AsyncMapMultiprocServer  # wherever your class is
+
+sock_prefix = f"/tmp/test_async_multiproc_map_{uuid.uuid4().hex}_"
+
+
+async def async_handler(keys, datum: Datum) -> Messages:
+    msg = f"payload:{datum.value.decode()} event_time:{datum.event_time} watermark:{datum.watermark}"
+    return Messages(Message(value=msg.encode(), keys=keys))
+
+class TestAsyncMapMultiprocServer(unittest.TestCase):
+    def setUp(self):
+        self.base_sock_path = sock_prefix
+        self.server = AsyncMapMultiprocServer(
+            mapper_instance=async_handler,
+            server_count=2,
+            sock_path=self.base_sock_path,
+            use_tcp=False,
+            server_info_file="/tmp/server_info"
+        )
+        self.process = Process(target=self.server.start)
+        self.process.start()
+
+        # Wait for both servers to bind
+        self.socket_paths = [f"{self.base_sock_path}{i}.sock" for i in range(2)]
+        for path in self.socket_paths:
+            for _ in range(10):
+                if os.path.exists(path):
+                    break
+                time.sleep(0.5)
+
+    def tearDown(self):
+        self.process.terminate()
+        self.process.join()
+        for path in self.socket_paths:
+            try:
+                os.remove(path)
+            except FileNotFoundError:
+                pass
+
+    def test_map_fn(self):
+        bind_address = f"unix://{self.socket_paths[0]}"
+        request = get_test_datums()
+        with grpc.insecure_channel(bind_address) as channel:
+            stub = map_pb2_grpc.MapStub(channel)
+            responses_iter = stub.MapFn(request_iterator=request_generator(request))
+            responses = []
+            # capture the output from the ReadFn generator and assert.
+            for r in responses_iter:
+                responses.append(r)
+
+            # 1 handshake + 3 data responses
+            self.assertEqual(4, len(responses))
+
+            self.assertTrue(responses[0].handshake.sot)
+
+            idx = 1
+            while idx < len(responses):
+                _id = "test-id-" + str(idx)
+                self.assertEqual(_id, responses[idx].id)
+                self.assertEqual(
+                    bytes(
+                        "payload:test_mock_message "
+                        "event_time:2022-09-12 16:00:00 watermark:2022-09-12 16:01:00",
+                        encoding="utf-8",
+                    ),
+                    responses[idx].results[0].value,
+                )
+                self.assertEqual(1, len(responses[idx].results))
+                idx += 1
+
+    def test_server_start(self):
+        for path in self.socket_paths:
+            self.assertTrue(os.path.exists(path), f"Server socket {path} was not created successfully")