Use @pytest.mark.parametrize to reduce test duplication

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

Author: BulkBeing

packages/pynumaflow/tests/accumulator/test_async_accumulator.py

@@ -409,26 +409,22 @@ class ExampleBadClass:
         AccumulatorAsyncServer(accumulator_instance=ExampleBadClass)
 
 
-def test_max_threads():
-    # max cap at 16
-    server = AccumulatorAsyncServer(accumulator_instance=ExampleClass, max_threads=32)
-    assert server.max_threads == 16
-
-    # use argument provided
-    server = AccumulatorAsyncServer(accumulator_instance=ExampleClass, max_threads=5)
-    assert server.max_threads == 5
-
-    # defaults to 4
-    server = AccumulatorAsyncServer(accumulator_instance=ExampleClass)
-    assert server.max_threads == 4
-
-    # zero threads
-    server = AccumulatorAsyncServer(ExampleClass, max_threads=0)
-    assert server.max_threads == 0
-
-    # negative threads
-    server = AccumulatorAsyncServer(ExampleClass, max_threads=-5)
-    assert server.max_threads == -5
+@pytest.mark.parametrize(
+    "max_threads_arg,expected",
+    [
+        (32, 16),  # max cap at 16
+        (5, 5),  # use argument provided
+        (None, 4),  # defaults to 4
+        (0, 0),  # zero threads
+        (-5, -5),  # negative threads
+    ],
+)
+def test_max_threads(max_threads_arg, expected):
+    kwargs = {"accumulator_instance": ExampleClass}
+    if max_threads_arg is not None:
+        kwargs["max_threads"] = max_threads_arg
+    server = AccumulatorAsyncServer(**kwargs)
+    assert server.max_threads == expected
 
 
 def test_server_info_file_path_handling():

packages/pynumaflow/tests/batchmap/test_async_batch_map.py

@@ -169,15 +169,17 @@ def test_is_ready(async_batch_map_server) -> None:
         assert response.ready
 
 
-def test_max_threads():
-    # max cap at 16
-    server = BatchMapAsyncServer(batch_mapper_instance=handler, max_threads=32)
-    assert server.max_threads == 16
-
-    # use argument provided
-    server = BatchMapAsyncServer(batch_mapper_instance=handler, max_threads=5)
-    assert server.max_threads == 5
-
-    # defaults to 4
-    server = BatchMapAsyncServer(batch_mapper_instance=handler)
-    assert server.max_threads == 4
+@pytest.mark.parametrize(
+    "max_threads_arg,expected",
+    [
+        (32, 16),  # max cap at 16
+        (5, 5),  # use argument provided
+        (None, 4),  # defaults to 4
+    ],
+)
+def test_max_threads(max_threads_arg, expected):
+    kwargs = {"batch_mapper_instance": handler}
+    if max_threads_arg is not None:
+        kwargs["max_threads"] = max_threads_arg
+    server = BatchMapAsyncServer(**kwargs)
+    assert server.max_threads == expected

packages/pynumaflow/tests/map/test_async_mapper.py

@@ -206,15 +206,17 @@ def test_invalid_input():
         MapAsyncServer()
 
 
-def test_max_threads():
-    # max cap at 16
-    server = MapAsyncServer(mapper_instance=async_map_handler, max_threads=32)
-    assert server.max_threads == 16
-
-    # use argument provided
-    server = MapAsyncServer(mapper_instance=async_map_handler, max_threads=5)
-    assert server.max_threads == 5
-
-    # defaults to 4
-    server = MapAsyncServer(mapper_instance=async_map_handler)
-    assert server.max_threads == 4
+@pytest.mark.parametrize(
+    "max_threads_arg,expected",
+    [
+        (32, 16),  # max cap at 16
+        (5, 5),  # use argument provided
+        (None, 4),  # defaults to 4
+    ],
+)
+def test_max_threads(max_threads_arg, expected):
+    kwargs = {"mapper_instance": async_map_handler}
+    if max_threads_arg is not None:
+        kwargs["max_threads"] = max_threads_arg
+    server = MapAsyncServer(**kwargs)
+    assert server.max_threads == expected

packages/pynumaflow/tests/map/test_multiproc_mapper.py

@@ -27,51 +27,41 @@ def _invoke_map_fn(test_server, timeout=1):
     )
 
 
-def test_multiproc_init():
-    my_server = MapMultiprocServer(mapper_instance=map_handler, server_count=3)
-    assert my_server._process_count == 3
-
-
-def test_multiproc_process_count():
-    default_val = os.cpu_count()
-    my_server = MapMultiprocServer(mapper_instance=map_handler)
-    assert my_server._process_count == default_val
-
-
-def test_max_process_count():
-    """Max process count is capped at 2 * os.cpu_count, irrespective of what the user
-    provides as input"""
-    default_val = os.cpu_count()
-    server = MapMultiprocServer(mapper_instance=map_handler, server_count=100)
-    assert server._process_count == default_val * 2
-
-
-def test_udf_map_err_handshake():
+@pytest.mark.parametrize(
+    "server_count,expected",
+    [
+        (3, 3),  # explicit count
+        (None, os.cpu_count()),  # default to cpu count
+        (100, os.cpu_count() * 2),  # max cap at 2 * cpu count
+    ],
+)
+def test_process_count(server_count, expected):
+    kwargs = {"mapper_instance": map_handler}
+    if server_count is not None:
+        kwargs["server_count"] = server_count
+    server = MapMultiprocServer(**kwargs)
+    assert server._process_count == expected
+
+
+@pytest.mark.parametrize(
+    "handshake,expected_msg",
+    [
+        (False, "MapFn: expected handshake as the first message"),
+        (True, "Something is fishy!"),
+    ],
+)
+def test_udf_map_error(handshake, expected_msg):
     my_server = MapMultiprocServer(mapper_instance=err_map_handler)
     services = {map_pb2.DESCRIPTOR.services_by_name["Map"]: my_server.servicer}
     test_server = server_from_dictionary(services, strict_real_time())
 
-    test_datums = get_test_datums(handshake=False)
-    method = _invoke_map_fn(test_server)
-    send_test_requests(method, test_datums)
-    drain_responses(method)
-
-    metadata, code, details = method.termination()
-    assert "MapFn: expected handshake as the first message" in details
-    assert code == StatusCode.INTERNAL
-
-
-def test_udf_map_err():
-    my_server = MapMultiprocServer(mapper_instance=err_map_handler)
-    services = {map_pb2.DESCRIPTOR.services_by_name["Map"]: my_server.servicer}
-    test_server = server_from_dictionary(services, strict_real_time())
-    test_datums = get_test_datums(handshake=True)
+    test_datums = get_test_datums(handshake=handshake)
     method = _invoke_map_fn(test_server)
     send_test_requests(method, test_datums)
     drain_responses(method)
 
     metadata, code, details = method.termination()
-    assert "Something is fishy!" in details
+    assert expected_msg in details
     assert code == StatusCode.INTERNAL
 
 

packages/pynumaflow/tests/map/test_sync_mapper.py

@@ -36,33 +36,25 @@ def test_init_with_args():
     assert my_servicer.max_message_size == 1024 * 1024 * 5
 
 
-def test_udf_map_err_handshake():
+@pytest.mark.parametrize(
+    "handshake,expected_msg",
+    [
+        (False, "MapFn: expected handshake as the first message"),
+        (True, "Something is fishy!"),
+    ],
+)
+def test_udf_map_error(handshake, expected_msg):
     my_server = MapServer(mapper_instance=err_map_handler)
     services = {map_pb2.DESCRIPTOR.services_by_name["Map"]: my_server.servicer}
     test_server = server_from_dictionary(services, strict_real_time())
 
-    test_datums = get_test_datums(handshake=False)
+    test_datums = get_test_datums(handshake=handshake)
     method = _invoke_map_fn(test_server)
     send_test_requests(method, test_datums)
     drain_responses(method)
 
     metadata, code, details = method.termination()
-    assert "MapFn: expected handshake as the first message" in details
-    assert code == StatusCode.INTERNAL
-
-
-def test_udf_map_error_response():
-    my_server = MapServer(mapper_instance=err_map_handler)
-    services = {map_pb2.DESCRIPTOR.services_by_name["Map"]: my_server.servicer}
-    test_server = server_from_dictionary(services, strict_real_time())
-
-    test_datums = get_test_datums(handshake=True)
-    method = _invoke_map_fn(test_server)
-    send_test_requests(method, test_datums)
-    drain_responses(method)
-
-    metadata, code, details = method.termination()
-    assert "Something is fishy!" in details
+    assert expected_msg in details
     assert code == StatusCode.INTERNAL
 
 
@@ -110,15 +102,17 @@ def test_invalid_input():
         MapServer()
 
 
-def test_max_threads():
-    # max cap at 16
-    server = MapServer(mapper_instance=map_handler, max_threads=32)
-    assert server.max_threads == 16
-
-    # use argument provided
-    server = MapServer(mapper_instance=map_handler, max_threads=5)
-    assert server.max_threads == 5
-
-    # defaults to 4
-    server = MapServer(mapper_instance=map_handler)
-    assert server.max_threads == 4
+@pytest.mark.parametrize(
+    "max_threads_arg,expected",
+    [
+        (32, 16),  # max cap at 16
+        (5, 5),  # use argument provided
+        (None, 4),  # defaults to 4
+    ],
+)
+def test_max_threads(max_threads_arg, expected):
+    kwargs = {"mapper_instance": map_handler}
+    if max_threads_arg is not None:
+        kwargs["max_threads"] = max_threads_arg
+    server = MapServer(**kwargs)
+    assert server.max_threads == expected

packages/pynumaflow/tests/mapstream/test_async_map_stream.py

@@ -153,15 +153,17 @@ def test_is_ready(async_map_stream_server):
         assert response.ready
 
 
-def test_max_threads():
-    # max cap at 16
-    server = MapStreamAsyncServer(map_stream_instance=async_map_stream_handler, max_threads=32)
-    assert server.max_threads == 16
-
-    # use argument provided
-    server = MapStreamAsyncServer(map_stream_instance=async_map_stream_handler, max_threads=5)
-    assert server.max_threads == 5
-
-    # defaults to 4
-    server = MapStreamAsyncServer(map_stream_instance=async_map_stream_handler)
-    assert server.max_threads == 4
+@pytest.mark.parametrize(
+    "max_threads_arg,expected",
+    [
+        (32, 16),  # max cap at 16
+        (5, 5),  # use argument provided
+        (None, 4),  # defaults to 4
+    ],
+)
+def test_max_threads(max_threads_arg, expected):
+    kwargs = {"map_stream_instance": async_map_stream_handler}
+    if max_threads_arg is not None:
+        kwargs["max_threads"] = max_threads_arg
+    server = MapStreamAsyncServer(**kwargs)
+    assert server.max_threads == expected

packages/pynumaflow/tests/reduce/test_async_reduce.py

@@ -248,15 +248,17 @@ class ExampleBadClass:
         ReduceAsyncServer(reducer_instance=ExampleBadClass)
 
 
-def test_max_threads():
-    # max cap at 16
-    server = ReduceAsyncServer(reducer_instance=ExampleClass, max_threads=32)
-    assert server.max_threads == 16
-
-    # use argument provided
-    server = ReduceAsyncServer(reducer_instance=ExampleClass, max_threads=5)
-    assert server.max_threads == 5
-
-    # defaults to 4
-    server = ReduceAsyncServer(reducer_instance=ExampleClass)
-    assert server.max_threads == 4
+@pytest.mark.parametrize(
+    "max_threads_arg,expected",
+    [
+        (32, 16),  # max cap at 16
+        (5, 5),  # use argument provided
+        (None, 4),  # defaults to 4
+    ],
+)
+def test_max_threads(max_threads_arg, expected):
+    kwargs = {"reducer_instance": ExampleClass}
+    if max_threads_arg is not None:
+        kwargs["max_threads"] = max_threads_arg
+    server = ReduceAsyncServer(**kwargs)
+    assert server.max_threads == expected

packages/pynumaflow/tests/reducestreamer/test_async_reduce.py

@@ -279,18 +279,20 @@ class ExampleBadClass:
         ReduceStreamAsyncServer(reduce_stream_instance=ExampleBadClass)
 
 
-def test_max_threads():
-    # max cap at 16
-    server = ReduceStreamAsyncServer(reduce_stream_instance=ExampleClass, max_threads=32)
-    assert server.max_threads == 16
-
-    # use argument provided
-    server = ReduceStreamAsyncServer(reduce_stream_instance=ExampleClass, max_threads=5)
-    assert server.max_threads == 5
-
-    # defaults to 4
-    server = ReduceStreamAsyncServer(reduce_stream_instance=ExampleClass)
-    assert server.max_threads == 4
+@pytest.mark.parametrize(
+    "max_threads_arg,expected",
+    [
+        (32, 16),  # max cap at 16
+        (5, 5),  # use argument provided
+        (None, 4),  # defaults to 4
+    ],
+)
+def test_max_threads(max_threads_arg, expected):
+    kwargs = {"reduce_stream_instance": ExampleClass}
+    if max_threads_arg is not None:
+        kwargs["max_threads"] = max_threads_arg
+    server = ReduceStreamAsyncServer(**kwargs)
+    assert server.max_threads == expected
 
 
 def test_start_shutdown_handler_without_callback():

packages/pynumaflow/tests/sideinput/test_side_input_server.py

@@ -108,15 +108,17 @@ def test_invalid_input():
         SideInputServer()
 
 
-def test_max_threads():
-    # max cap at 16
-    server = SideInputServer(retrieve_side_input_handler, max_threads=32)
-    assert server.max_threads == 16
-
-    # use argument provided
-    server = SideInputServer(retrieve_side_input_handler, max_threads=5)
-    assert server.max_threads == 5
-
-    # defaults to 4
-    server = SideInputServer(retrieve_side_input_handler)
-    assert server.max_threads == 4
+@pytest.mark.parametrize(
+    "max_threads_arg,expected",
+    [
+        (32, 16),  # max cap at 16
+        (5, 5),  # use argument provided
+        (None, 4),  # defaults to 4
+    ],
+)
+def test_max_threads(max_threads_arg, expected):
+    kwargs = {"side_input_instance": retrieve_side_input_handler}
+    if max_threads_arg is not None:
+        kwargs["max_threads"] = max_threads_arg
+    server = SideInputServer(**kwargs)
+    assert server.max_threads == expected