conftest.py with shared fixtures

Signed-off-by: Sreekanth <prsreekanth920@gmail.com>

Author: BulkBeing

packages/pynumaflow/tests/conftest.py

@@ -0,0 +1,64 @@
+"""
+Root conftest.py — shared pytest fixtures and helpers for all test modules.
+
+Provides helpers for common gRPC testing patterns that are duplicated across
+sync, multiproc, and async test files.
+"""
+
+
+def collect_responses(method):
+    """Collect all responses from a grpc_testing stream method until exhausted.
+
+    Replaces the repeated pattern:
+        responses = []
+        while True:
+            try:
+                resp = method.take_response()
+                responses.append(resp)
+            except ValueError as err:
+                if "No more responses!" in err.__str__():
+                    break
+
+    Returns a list of response protos.
+    """
+    responses = []
+    while True:
+        try:
+            resp = method.take_response()
+            responses.append(resp)
+        except ValueError as err:
+            if "No more responses!" in str(err):
+                break
+    return responses
+
+
+def drain_responses(method):
+    """Drain all responses from a grpc_testing stream method, discarding them.
+
+    Replaces the repeated pattern:
+        while True:
+            try:
+                method.take_response()
+            except ValueError:
+                break
+
+    Useful in shutdown tests where we only care about termination status.
+    """
+    while True:
+        try:
+            method.take_response()
+        except ValueError:
+            break
+
+
+def send_test_requests(method, datums):
+    """Send a list of test datums to a grpc_testing stream method and close.
+
+    Replaces the repeated pattern:
+        for d in test_datums:
+            method.send_request(d)
+        method.requests_closed()
+    """
+    for d in datums:
+        method.send_request(d)
+    method.requests_closed()

packages/pynumaflow/tests/map/test_multiproc_map_shutdown.py

@@ -14,6 +14,7 @@
 
 from pynumaflow.mapper import MapMultiprocServer
 from pynumaflow.proto.mapper import map_pb2
+from tests.conftest import drain_responses, send_test_requests
 from tests.map.utils import map_handler, err_map_handler, get_test_datums
 
 
@@ -33,15 +34,8 @@ def test_shutdown_event_set_on_handler_error():
         timeout=2,
     )
 
-    for d in test_datums:
-        method.send_request(d)
-    method.requests_closed()
-
-    while True:
-        try:
-            method.take_response()
-        except ValueError:
-            break
+    send_test_requests(method, test_datums)
+    drain_responses(method)
 
     _, code, _ = method.termination()
     assert code == StatusCode.INTERNAL
@@ -65,15 +59,8 @@ def test_shutdown_event_set_on_handshake_error():
         timeout=1,
     )
 
-    for d in test_datums:
-        method.send_request(d)
-    method.requests_closed()
-
-    while True:
-        try:
-            method.take_response()
-        except ValueError:
-            break
+    send_test_requests(method, test_datums)
+    drain_responses(method)
 
     _, code, details = method.termination()
     assert code == StatusCode.INTERNAL

packages/pynumaflow/tests/map/test_multiproc_mapper.py

@@ -10,6 +10,7 @@
 from pynumaflow.mapper import MapMultiprocServer
 from pynumaflow.proto.mapper import map_pb2
 from tests.map.utils import map_handler, err_map_handler, get_test_datums
+from tests.conftest import collect_responses, drain_responses, send_test_requests
 from tests.testing_utils import (
     mock_terminate_on_stop,
 )
@@ -50,18 +51,8 @@ def test_udf_map_err_handshake(self):
             invocation_metadata={},
             timeout=1,
         )
-        for x in test_datums:
-            method.send_request(x)
-        method.requests_closed()
-
-        responses = []
-        while True:
-            try:
-                resp = method.take_response()
-                responses.append(resp)
-            except ValueError as err:
-                if "No more responses!" in err.__str__():
-                    break
+        send_test_requests(method, test_datums)
+        drain_responses(method)
 
         metadata, code, details = method.termination()
         self.assertTrue("MapFn: expected handshake as the first message" in details)
@@ -77,18 +68,8 @@ def test_udf_map_err(self):
             invocation_metadata={},
             timeout=1,
         )
-        for x in test_datums:
-            method.send_request(x)
-        method.requests_closed()
-
-        responses = []
-        while True:
-            try:
-                resp = method.take_response()
-                responses.append(resp)
-            except ValueError as err:
-                if "No more responses!" in err.__str__():
-                    break
+        send_test_requests(method, test_datums)
+        drain_responses(method)
 
         metadata, code, details = method.termination()
         self.assertTrue("Something is fishy!" in details)
@@ -116,18 +97,8 @@ def test_map_forward_message(self):
             invocation_metadata={},
             timeout=1,
         )
-        for x in test_datums:
-            method.send_request(x)
-        method.requests_closed()
-
-        responses = []
-        while True:
-            try:
-                resp = method.take_response()
-                responses.append(resp)
-            except ValueError as err:
-                if "No more responses!" in err.__str__():
-                    break
+        send_test_requests(method, test_datums)
+        responses = collect_responses(method)
 
         metadata, code, details = method.termination()
 

packages/pynumaflow/tests/map/test_sync_map_shutdown.py

@@ -15,6 +15,7 @@
 
 from pynumaflow.mapper._servicer._sync_servicer import SyncMapServicer
 from pynumaflow.proto.mapper import map_pb2
+from tests.conftest import drain_responses, send_test_requests
 from tests.map.utils import map_handler, err_map_handler, get_test_datums
 
 
@@ -33,15 +34,8 @@ def test_shutdown_event_set_on_handler_error():
         timeout=2,
     )
 
-    for d in test_datums:
-        method.send_request(d)
-    method.requests_closed()
-
-    while True:
-        try:
-            method.take_response()
-        except ValueError:
-            break
+    send_test_requests(method, test_datums)
+    drain_responses(method)
 
     _, code, _ = method.termination()
     assert code == StatusCode.INTERNAL
@@ -65,15 +59,8 @@ def test_shutdown_event_set_on_handshake_error():
         timeout=1,
     )
 
-    for d in test_datums:
-        method.send_request(d)
-    method.requests_closed()
-
-    while True:
-        try:
-            method.take_response()
-        except ValueError:
-            break
+    send_test_requests(method, test_datums)
+    drain_responses(method)
 
     _, code, details = method.termination()
     assert code == StatusCode.INTERNAL

packages/pynumaflow/tests/map/test_sync_mapper.py

@@ -9,6 +9,7 @@
 from pynumaflow.mapper import MapServer
 from pynumaflow.proto.mapper import map_pb2
 from tests.map.utils import map_handler, err_map_handler, ExampleMap, get_test_datums
+from tests.conftest import collect_responses, drain_responses, send_test_requests
 from tests.testing_utils import (
     mock_terminate_on_stop,
 )
@@ -44,18 +45,8 @@ def test_udf_map_err_handshake(self):
             invocation_metadata={},
             timeout=1,
         )
-        for x in test_datums:
-            method.send_request(x)
-        method.requests_closed()
-
-        responses = []
-        while True:
-            try:
-                resp = method.take_response()
-                responses.append(resp)
-            except ValueError as err:
-                if "No more responses!" in err.__str__():
-                    break
+        send_test_requests(method, test_datums)
+        drain_responses(method)
 
         metadata, code, details = method.termination()
         self.assertTrue("MapFn: expected handshake as the first message" in details)
@@ -72,18 +63,8 @@ def test_udf_map_error_response(self):
             invocation_metadata={},
             timeout=1,
         )
-        for x in test_datums:
-            method.send_request(x)
-        method.requests_closed()
-
-        responses = []
-        while True:
-            try:
-                resp = method.take_response()
-                responses.append(resp)
-            except ValueError as err:
-                if "No more responses!" in err.__str__():
-                    break
+        send_test_requests(method, test_datums)
+        drain_responses(method)
 
         metadata, code, details = method.termination()
         self.assertTrue("Something is fishy!" in details)
@@ -111,18 +92,8 @@ def test_map_forward_message(self):
             invocation_metadata={},
             timeout=1,
         )
-        for x in test_datums:
-            method.send_request(x)
-        method.requests_closed()
-
-        responses = []
-        while True:
-            try:
-                resp = method.take_response()
-                responses.append(resp)
-            except ValueError as err:
-                if "No more responses!" in err.__str__():
-                    break
+        send_test_requests(method, test_datums)
+        responses = collect_responses(method)
 
         metadata, code, details = method.termination()
         # 1 handshake + 3 data responses

packages/pynumaflow/tests/sink/test_server.py

@@ -22,6 +22,7 @@
 from pynumaflow.proto.sinker import sink_pb2
 from pynumaflow.sinker import Responses, Datum, Response, SinkServer, Message, UserMetadata
 from pynumaflow.sinker.servicer.sync_servicer import SyncSinkServicer
+from tests.conftest import collect_responses, drain_responses, send_test_requests
 
 
 def mockenv(**envvars):
@@ -201,18 +202,8 @@ def test_udsink_err(err_sink_test_server):
         timeout=1,
     )
 
-    for d in test_datums:
-        method.send_request(d)
-    method.requests_closed()
-
-    responses = []
-    while True:
-        try:
-            resp = method.take_response()
-            responses.append(resp)
-        except ValueError as err:
-            if "No more responses!" in str(err):
-                break
+    send_test_requests(method, test_datums)
+    drain_responses(method)
 
     metadata, code, details = method.termination()
     assert code == StatusCode.INTERNAL
@@ -267,18 +258,8 @@ def test_forward_message(sink_test_server):
         invocation_metadata={},
         timeout=1,
     )
-    for x in test_datums:
-        method.send_request(x)
-    method.requests_closed()
-
-    responses = []
-    while True:
-        try:
-            resp = method.take_response()
-            responses.append(resp)
-        except ValueError as err:
-            if "No more responses!" in str(err):
-                break
+    send_test_requests(method, test_datums)
+    responses = collect_responses(method)
 
     # 1 handshake +  1 data messages + 1 EOT
     assert len(responses) == 3
@@ -370,15 +351,8 @@ def test_shutdown_event_set_on_handler_error():
         timeout=2,
     )
 
-    for d in test_datums:
-        method.send_request(d)
-    method.requests_closed()
-
-    while True:
-        try:
-            method.take_response()
-        except ValueError:
-            break
+    send_test_requests(method, test_datums)
+    drain_responses(method)
 
     _, code, _ = method.termination()
     assert code == StatusCode.INTERNAL