more tests, remove unused variables

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

Author: BulkBeing

packages/pynumaflow/pynumaflow/mapper/_servicer/_sync_servicer.py

@@ -21,10 +21,8 @@ class SyncMapServicer(map_pb2_grpc.MapServicer):
     Provides the functionality for the required rpc methods.
     """
 
-    def __init__(self, handler: MapSyncCallable, multiproc: bool = False):
+    def __init__(self, handler: MapSyncCallable):
         self.__map_handler: MapSyncCallable = handler
-        # This indicates whether the grpc server attached is multiproc or not
-        self.multiproc = multiproc
         # create a thread pool for executing UDF code
         self.executor = ThreadPoolExecutor(max_workers=NUM_THREADS_DEFAULT)
         # Graceful shutdown: when set, a watcher thread in _run_server() calls

packages/pynumaflow/pynumaflow/mapper/multiproc_server.py

@@ -106,7 +106,7 @@ def handler(self, keys: list[str], datum: Datum) -> Messages:
         # Setting the max value to 2 * CPU count
         # Used for multiproc server
         self._process_count = min(server_count, 2 * _PROCESS_COUNT)
-        self.servicer = SyncMapServicer(handler=mapper_instance, multiproc=True)
+        self.servicer = SyncMapServicer(handler=mapper_instance)
 
         # Shared event across all worker processes for coordinated shutdown.
         # When any worker's servicer sets this event, all workers' watcher

packages/pynumaflow/pynumaflow/sourcetransformer/multiproc_server.py

@@ -131,7 +131,7 @@ def my_handler(keys: list[str], datum: Datum) -> Messages:
         # Setting the max value to 2 * CPU count
         # Used for multiproc server
         self._process_count = min(server_count, 2 * _PROCESS_COUNT)
-        self.servicer = SourceTransformServicer(handler=source_transform_instance, multiproc=True)
+        self.servicer = SourceTransformServicer(handler=source_transform_instance)
 
         # Shared event across all worker processes for coordinated shutdown.
         # When any worker's servicer sets this event, all workers' watcher

packages/pynumaflow/pynumaflow/sourcetransformer/servicer/_servicer.py

@@ -41,10 +41,8 @@ class SourceTransformServicer(transform_pb2_grpc.SourceTransformServicer):
     Provides the functionality for the required rpc methods.
     """
 
-    def __init__(self, handler: SourceTransformCallable, multiproc: bool = False):
+    def __init__(self, handler: SourceTransformCallable):
         self.__transform_handler: SourceTransformCallable = handler
-        # This indicates whether the grpc server attached is multiproc or not
-        self.multiproc = multiproc
         # create a thread pool for executing UDF code
         self.executor = ThreadPoolExecutor(max_workers=NUM_THREADS_DEFAULT)
         # Graceful shutdown: when set, a watcher thread in _run_server() calls

packages/pynumaflow/tests/map/test_multiproc_map_shutdown.py

@@ -0,0 +1,119 @@
+"""
+Shutdown-event tests for the multiproc Map servicer.
+
+These tests verify that the SyncMapServicer (as used by MapMultiprocServer)
+correctly sets shutdown_event on error, enabling coordinated graceful shutdown
+across all worker processes via the shared multiprocessing.Event.
+"""
+
+from unittest import mock
+
+import grpc
+from grpc import StatusCode
+from grpc_testing import server_from_dictionary, strict_real_time
+
+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
+
+
+def test_shutdown_event_set_on_handler_error():
+    """When the UDF handler raises, the servicer must signal the shutdown event."""
+    server = MapMultiprocServer(mapper_instance=err_map_handler)
+    servicer = server.servicer
+
+    services = {map_pb2.DESCRIPTOR.services_by_name["Map"]: servicer}
+    test_server = server_from_dictionary(services, strict_real_time())
+
+    test_datums = get_test_datums(handshake=True)
+
+    method = test_server.invoke_stream_stream(
+        method_descriptor=(map_pb2.DESCRIPTOR.services_by_name["Map"].methods_by_name["MapFn"]),
+        invocation_metadata={},
+        timeout=2,
+    )
+
+    for d in test_datums:
+        method.send_request(d)
+    method.requests_closed()
+
+    while True:
+        try:
+            method.take_response()
+        except ValueError:
+            break
+
+    _, code, _ = method.termination()
+    assert code == StatusCode.INTERNAL
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is not None
+
+
+def test_shutdown_event_set_on_handshake_error():
+    """Missing handshake must also signal the shutdown event."""
+    server = MapMultiprocServer(mapper_instance=map_handler)
+    servicer = server.servicer
+
+    services = {map_pb2.DESCRIPTOR.services_by_name["Map"]: servicer}
+    test_server = server_from_dictionary(services, strict_real_time())
+
+    test_datums = get_test_datums(handshake=False)
+
+    method = test_server.invoke_stream_stream(
+        method_descriptor=(map_pb2.DESCRIPTOR.services_by_name["Map"].methods_by_name["MapFn"]),
+        invocation_metadata={},
+        timeout=1,
+    )
+
+    for d in test_datums:
+        method.send_request(d)
+    method.requests_closed()
+
+    while True:
+        try:
+            method.take_response()
+        except ValueError:
+            break
+
+    _, code, details = method.termination()
+    assert code == StatusCode.INTERNAL
+    assert "MapFn: expected handshake as the first message" in details
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is not None
+
+
+def test_shutdown_event_set_on_stream_close_before_handshake():
+    """grpc.RpcError on the first read (before handshake): shutdown_event set,
+    result_queue is None so close is skipped."""
+    server = MapMultiprocServer(mapper_instance=map_handler)
+    servicer = server.servicer
+
+    def _cancelled_iter():
+        raise grpc.RpcError()
+        yield  # make it a generator
+
+    responses = list(servicer.MapFn(_cancelled_iter(), mock.MagicMock()))
+
+    assert responses == []
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is None
+
+
+def test_shutdown_event_set_on_stream_close_mid_processing():
+    """grpc.RpcError mid-processing: result_queue is closed (unblocking the handler
+    thread) and shutdown_event is set."""
+    server = MapMultiprocServer(mapper_instance=map_handler)
+    servicer = server.servicer
+
+    test_datums = get_test_datums(handshake=True)
+
+    def _cancelled_iter():
+        yield test_datums[0]  # handshake
+        yield test_datums[1]  # first data message
+        raise grpc.RpcError()
+
+    responses = list(servicer.MapFn(_cancelled_iter(), mock.MagicMock()))
+
+    assert responses[0].handshake.sot
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is None

packages/pynumaflow/tests/sideinput/test_shutdown.py

@@ -0,0 +1,72 @@
+"""
+Shutdown-event tests for the SideInput servicer.
+
+Verifies that the servicer sets shutdown_event and captures the error when the
+UDF handler raises, enabling graceful server stop via the watcher thread in
+_run_server() instead of a hard process kill.
+"""
+
+from grpc import StatusCode
+from grpc_testing import server_from_dictionary, strict_real_time
+from google.protobuf import empty_pb2 as _empty_pb2
+
+from pynumaflow.sideinput.servicer.servicer import SideInputServicer
+from pynumaflow.proto.sideinput import sideinput_pb2
+
+
+def _ok_handler():
+    from pynumaflow.sideinput import Response
+
+    return Response.broadcast_message(b"test")
+
+
+def _err_handler():
+    raise RuntimeError("Something is fishy!")
+
+
+def test_shutdown_event_set_on_handler_error():
+    """When the UDF handler raises, the servicer must signal the shutdown event."""
+    servicer = SideInputServicer(handler=_err_handler)
+
+    services = {sideinput_pb2.DESCRIPTOR.services_by_name["SideInput"]: servicer}
+    test_server = server_from_dictionary(services, strict_real_time())
+
+    method = test_server.invoke_unary_unary(
+        method_descriptor=(
+            sideinput_pb2.DESCRIPTOR.services_by_name["SideInput"].methods_by_name[
+                "RetrieveSideInput"
+            ]
+        ),
+        invocation_metadata={},
+        request=_empty_pb2.Empty(),
+        timeout=1,
+    )
+
+    _, _, code, _ = method.termination()
+    assert code == StatusCode.INTERNAL
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is not None
+
+
+def test_shutdown_event_not_set_on_success():
+    """On a successful call, shutdown_event must remain unset."""
+    servicer = SideInputServicer(handler=_ok_handler)
+
+    services = {sideinput_pb2.DESCRIPTOR.services_by_name["SideInput"]: servicer}
+    test_server = server_from_dictionary(services, strict_real_time())
+
+    method = test_server.invoke_unary_unary(
+        method_descriptor=(
+            sideinput_pb2.DESCRIPTOR.services_by_name["SideInput"].methods_by_name[
+                "RetrieveSideInput"
+            ]
+        ),
+        invocation_metadata={},
+        request=_empty_pb2.Empty(),
+        timeout=1,
+    )
+
+    _, _, code, _ = method.termination()
+    assert code == StatusCode.OK
+    assert not servicer.shutdown_event.is_set()
+    assert servicer.error is None

packages/pynumaflow/tests/sourcetransform/test_multiproc_shutdown.py

@@ -0,0 +1,128 @@
+"""
+Shutdown-event tests for the multiproc SourceTransform servicer.
+
+These tests verify that the SourceTransformServicer (as used by
+SourceTransformMultiProcServer) correctly sets shutdown_event on error,
+enabling coordinated graceful shutdown across all worker processes via
+the shared multiprocessing.Event.
+"""
+
+from unittest import mock
+
+import grpc
+from grpc import StatusCode
+from grpc_testing import server_from_dictionary, strict_real_time
+
+from pynumaflow.sourcetransformer.multiproc_server import SourceTransformMultiProcServer
+from pynumaflow.proto.sourcetransformer import transform_pb2
+from tests.sourcetransform.utils import transform_handler, err_transform_handler, get_test_datums
+
+
+def test_shutdown_event_set_on_handler_error():
+    """When the UDF handler raises, the servicer must signal the shutdown event."""
+    server = SourceTransformMultiProcServer(source_transform_instance=err_transform_handler)
+    servicer = server.servicer
+
+    services = {transform_pb2.DESCRIPTOR.services_by_name["SourceTransform"]: servicer}
+    test_server = server_from_dictionary(services, strict_real_time())
+
+    test_datums = get_test_datums(handshake=True)
+
+    method = test_server.invoke_stream_stream(
+        method_descriptor=(
+            transform_pb2.DESCRIPTOR.services_by_name["SourceTransform"].methods_by_name[
+                "SourceTransformFn"
+            ]
+        ),
+        invocation_metadata={},
+        timeout=2,
+    )
+
+    for d in test_datums:
+        method.send_request(d)
+    method.requests_closed()
+
+    while True:
+        try:
+            method.take_response()
+        except ValueError:
+            break
+
+    _, code, _ = method.termination()
+    assert code == StatusCode.INTERNAL
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is not None
+
+
+def test_shutdown_event_set_on_handshake_error():
+    """Missing handshake must also signal the shutdown event."""
+    server = SourceTransformMultiProcServer(source_transform_instance=transform_handler)
+    servicer = server.servicer
+
+    services = {transform_pb2.DESCRIPTOR.services_by_name["SourceTransform"]: servicer}
+    test_server = server_from_dictionary(services, strict_real_time())
+
+    test_datums = get_test_datums(handshake=False)
+
+    method = test_server.invoke_stream_stream(
+        method_descriptor=(
+            transform_pb2.DESCRIPTOR.services_by_name["SourceTransform"].methods_by_name[
+                "SourceTransformFn"
+            ]
+        ),
+        invocation_metadata={},
+        timeout=1,
+    )
+
+    for d in test_datums:
+        method.send_request(d)
+    method.requests_closed()
+
+    while True:
+        try:
+            method.take_response()
+        except ValueError:
+            break
+
+    _, code, details = method.termination()
+    assert code == StatusCode.INTERNAL
+    assert "SourceTransformFn: expected handshake message" in details
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is not None
+
+
+def test_shutdown_event_set_on_stream_close_before_handshake():
+    """grpc.RpcError on the first read (before handshake): shutdown_event set,
+    result_queue is None so close is skipped."""
+    server = SourceTransformMultiProcServer(source_transform_instance=transform_handler)
+    servicer = server.servicer
+
+    def _cancelled_iter():
+        raise grpc.RpcError()
+        yield  # make it a generator
+
+    responses = list(servicer.SourceTransformFn(_cancelled_iter(), mock.MagicMock()))
+
+    assert responses == []
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is None
+
+
+def test_shutdown_event_set_on_stream_close_mid_processing():
+    """grpc.RpcError mid-processing: result_queue is closed (unblocking the handler
+    thread) and shutdown_event is set."""
+    server = SourceTransformMultiProcServer(source_transform_instance=transform_handler)
+    servicer = server.servicer
+
+    test_datums = get_test_datums(handshake=True)
+
+    def _cancelled_iter():
+        yield test_datums[0]  # handshake
+        yield test_datums[1]  # first data message
+        raise grpc.RpcError()
+
+    responses = list(servicer.SourceTransformFn(_cancelled_iter(), mock.MagicMock()))
+
+    assert responses[0].handshake.sot
+    assert servicer.shutdown_event.is_set()
+    assert servicer.error is None