add unit and integration tests

Signed-off-by: Abdullah Yildirim <a3hadi@protonmail.com>

Author: ayildirim21

packages/pynumaflow/tests/batchmap/test_async_batch_map.py

@@ -13,7 +13,9 @@
     BatchResponses,
     BatchResponse,
     BatchMapAsyncServer,
+    NackOptions,
 )
+from pynumaflow._constants import NACK
 from pynumaflow.proto.mapper import map_pb2_grpc
 from tests.batchmap.utils import request_generator
 from tests.conftest import create_async_loop, start_async_server, teardown_async_server
@@ -23,6 +25,9 @@
 LOGGER = setup_logging(__name__)
 
 listen_addr = "unix:///tmp/batch_map.sock"
+nack_listen_addr = "unix:///tmp/batch_map_nack.sock"
+
+NACK_TEST_OPTIONS = NackOptions(delay=1000, max_deliveries=3, reason="retry")
 
 
 class ExampleClass(BatchMapper):
@@ -137,6 +142,53 @@ def test_batch_map(batch_map_stub) -> None:
     assert 11 == len(data_resp)
 
 
+async def nack_handler(datums: AsyncIterable[Datum]) -> BatchResponses:
+    batch_responses = BatchResponses()
+    async for datum in datums:
+        batch_response = BatchResponse.from_id(datum.id)
+        batch_response.append(Message.to_nack(NACK_TEST_OPTIONS))
+        batch_responses.append(batch_response)
+    return batch_responses
+
+
+async def start_nack_server(udfs):
+    server = grpc.aio.server()
+    map_pb2_grpc.add_MapServicer_to_server(udfs, server)
+    server.add_insecure_port(nack_listen_addr)
+    logging.info("Starting nack server on %s", nack_listen_addr)
+    await server.start()
+    return server, nack_listen_addr
+
+
+@pytest.fixture(scope="module")
+def async_nack_batch_map_server():
+    loop = create_async_loop()
+    udfs = BatchMapAsyncServer(nack_handler).servicer
+    server = start_async_server(loop, start_nack_server(udfs))
+    yield loop
+    teardown_async_server(loop, server)
+
+
+def test_batch_map_nack(async_nack_batch_map_server) -> None:
+    with grpc.insecure_channel(nack_listen_addr) as channel:
+        stub = map_pb2_grpc.MapStub(channel)
+        generator_response = stub.MapFn(request_iterator=request_generator(count=5, session=1))
+
+        handshake = next(generator_response)
+        assert handshake.handshake.sot
+
+        data_resp = [r for r in generator_response]
+
+    # exclude the trailing EOT response
+    for resp in data_resp[:-1]:
+        result = resp.results[0]
+        assert NACK in result.tags
+        assert result.nack_options.delay == NACK_TEST_OPTIONS.delay
+        assert result.nack_options.max_deliveries == NACK_TEST_OPTIONS.max_deliveries
+        assert result.nack_options.reason == NACK_TEST_OPTIONS.reason
+    assert data_resp[-1].status.eot is True
+
+
 def test_is_ready(async_batch_map_server) -> None:
     with grpc.insecure_channel(listen_addr) as channel:
         stub = map_pb2_grpc.MapStub(channel)

packages/pynumaflow/tests/batchmap/test_messages.py

@@ -1,6 +1,7 @@
 import pytest
 
-from pynumaflow.batchmapper import Message, DROP, BatchResponse, BatchResponses
+from pynumaflow.batchmapper import Message, DROP, BatchResponse, BatchResponses, NackOptions
+from pynumaflow._constants import NACK
 from tests.batchmap.test_datatypes import TEST_ID
 from tests.testing_utils import mock_message
 
@@ -10,6 +11,27 @@ def _mock_message_object():
     return Message(value=value)
 
 
+def test_message_to_nack():
+    opts = NackOptions(max_deliveries=2)
+    msg = Message.to_nack(opts)
+    assert type(msg) is Message
+    assert msg.keys == []
+    assert msg.value == b""
+    assert msg.tags == [NACK]
+    assert msg.nack_options == opts
+
+
+def test_message_to_nack_default_opts():
+    msg = Message.to_nack()
+    assert msg.tags == [NACK]
+    assert msg.nack_options is None
+
+
+def test_message_default_nack_options():
+    msg = Message(mock_message())
+    assert msg.nack_options is None
+
+
 def test_batch_responses_init():
     batch_responses = BatchResponses()
     batch_response1 = BatchResponse.from_id(TEST_ID)

packages/pynumaflow/tests/map/test_async_mapper.py

@@ -16,7 +16,8 @@
 from pynumaflow.proto.common import metadata_pb2
 from pynumaflow.proto.mapper import map_pb2, map_pb2_grpc
 from tests.conftest import create_async_loop, start_async_server, teardown_async_server
-from tests.map.utils import get_test_datums
+from tests.map.utils import get_test_datums, async_nack_map_handler, NACK_TEST_OPTIONS
+from pynumaflow._constants import NACK
 
 pytestmark = pytest.mark.integration
 
@@ -26,6 +27,7 @@
 raise_error_from_map = False
 
 SOCK_PATH = "unix:///tmp/async_map.sock"
+NACK_SOCK_PATH = "unix:///tmp/async_map_nack.sock"
 
 
 def request_generator(req):
@@ -106,6 +108,54 @@ def test_run_server(async_map_server):
         LOGGER.info("Successfully validated the server")
 
 
+async def _start_nack_server(udfs):
+    _server_options = [
+        ("grpc.max_send_message_length", MAX_MESSAGE_SIZE),
+        ("grpc.max_receive_message_length", MAX_MESSAGE_SIZE),
+    ]
+    server = grpc.aio.server(options=_server_options)
+    map_pb2_grpc.add_MapServicer_to_server(udfs, server)
+    server.add_insecure_port(NACK_SOCK_PATH)
+    logging.info("Starting nack server on %s", NACK_SOCK_PATH)
+    await server.start()
+    return server, NACK_SOCK_PATH
+
+
+@pytest.fixture(scope="module")
+def async_nack_map_server():
+    """Module-scoped fixture: async gRPC map server whose handler nacks every message."""
+    loop = create_async_loop()
+
+    server_obj = MapAsyncServer(mapper_instance=async_nack_map_handler)
+    udfs = server_obj.servicer
+    server = start_async_server(loop, _start_nack_server(udfs))
+
+    yield loop
+
+    teardown_async_server(loop, server)
+
+
+def test_map_nack(async_nack_map_server):
+    with grpc.insecure_channel(NACK_SOCK_PATH) as channel:
+        stub = map_pb2_grpc.MapStub(channel)
+        request = get_test_datums()
+        generator_response = stub.MapFn(request_iterator=request_generator(request))
+
+        responses = list(generator_response)
+
+    # 1 handshake + 3 data responses
+    assert len(responses) == 4
+    assert responses[0].handshake.sot
+
+    for resp in responses[1:]:
+        assert len(resp.results) == 1
+        result = resp.results[0]
+        assert NACK in result.tags
+        assert result.nack_options.delay == NACK_TEST_OPTIONS.delay
+        assert result.nack_options.max_deliveries == NACK_TEST_OPTIONS.max_deliveries
+        assert result.nack_options.reason == NACK_TEST_OPTIONS.reason
+
+
 def test_map(map_stub):
     request = get_test_datums()
     generator_response: Iterator[map_pb2.MapResponse] = map_stub.MapFn(

packages/pynumaflow/tests/map/test_messages.py

@@ -1,6 +1,7 @@
 import pytest
 
-from pynumaflow.mapper import Messages, Message, DROP, Mapper, Datum
+from pynumaflow.mapper import Messages, Message, DROP, Mapper, Datum, NackOptions
+from pynumaflow._constants import NACK
 from tests.testing_utils import mock_message
 
 
@@ -35,6 +36,27 @@ def test_message_to_drop():
     assert mock_obj["Tags"] == msg.tags
 
 
+def test_message_to_nack():
+    opts = NackOptions(max_deliveries=2)
+    msg = Message.to_nack(opts)
+    assert type(msg) is Message
+    assert msg.keys == []
+    assert msg.value == b""
+    assert msg.tags == [NACK]
+    assert msg.nack_options == opts
+
+
+def test_message_to_nack_default_opts():
+    msg = Message.to_nack()
+    assert msg.tags == [NACK]
+    assert msg.nack_options is None
+
+
+def test_message_default_nack_options():
+    msg = Message(mock_message())
+    assert msg.nack_options is None
+
+
 def test_message_to():
     mock_obj = {"Keys": ["__KEY__"], "Value": mock_message(), "Tags": ["__TAG__"]}
     msg = Message(value=mock_obj["Value"], keys=mock_obj["Keys"], tags=mock_obj["Tags"])

packages/pynumaflow/tests/map/test_sync_mapper.py

@@ -5,7 +5,15 @@
 
 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.map.utils import (
+    map_handler,
+    err_map_handler,
+    ExampleMap,
+    get_test_datums,
+    nack_map_handler,
+    NACK_TEST_OPTIONS,
+)
+from pynumaflow._constants import NACK
 from tests.conftest import collect_responses, drain_responses, send_test_requests
 
 
@@ -97,6 +105,32 @@ def test_map_forward_message(map_test_server):
     assert code == StatusCode.OK
 
 
+def test_map_nack():
+    my_server = MapServer(mapper_instance=nack_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)
+    responses = collect_responses(method)
+
+    metadata, code, details = method.termination()
+    # 1 handshake + 3 data responses
+    assert len(responses) == 4
+    assert responses[0].handshake.sot
+
+    for resp in responses[1:]:
+        assert len(resp.results) == 1
+        result = resp.results[0]
+        # nack messages carry the NACK tag and the serialized nack options
+        assert NACK in result.tags
+        assert result.nack_options.delay == NACK_TEST_OPTIONS.delay
+        assert result.nack_options.max_deliveries == NACK_TEST_OPTIONS.max_deliveries
+        assert result.nack_options.reason == NACK_TEST_OPTIONS.reason
+    assert code == StatusCode.OK
+
+
 def test_invalid_input():
     with pytest.raises(TypeError):
         MapServer()

packages/pynumaflow/tests/map/utils.py

@@ -1,10 +1,22 @@
 from pynumaflow.proto.common import metadata_pb2
 from pynumaflow.proto.mapper import map_pb2
 
-from pynumaflow.mapper import Datum, Messages, Message, Mapper
+from pynumaflow.mapper import Datum, Messages, Message, Mapper, NackOptions
 from tests.testing_utils import mock_message, mock_headers, get_time_args
 
 
+# NackOptions used by the nack handlers below; also asserted by the round-trip tests.
+NACK_TEST_OPTIONS = NackOptions(delay=1000, max_deliveries=3, reason="retry")
+
+
+def nack_map_handler(keys: list[str], datum: Datum) -> Messages:
+    return Messages(Message.to_nack(NACK_TEST_OPTIONS))
+
+
+async def async_nack_map_handler(keys: list[str], datum: Datum) -> Messages:
+    return Messages(Message.to_nack(NACK_TEST_OPTIONS))
+
+
 async def async_map_error_fn(keys: list[str], datum: Datum) -> Messages:
     raise ValueError("error invoking map")
 

packages/pynumaflow/tests/mapstream/test_async_map_stream.py

@@ -10,7 +10,9 @@
     Message,
     Datum,
     MapStreamAsyncServer,
+    NackOptions,
 )
+from pynumaflow._constants import NACK
 from pynumaflow.proto.mapper import map_pb2_grpc
 from tests.mapstream.utils import request_generator
 from tests.conftest import create_async_loop, start_async_server, teardown_async_server
@@ -25,6 +27,9 @@
 raise_error_from_map = False
 
 SOCK_PATH = "unix:///tmp/async_map_stream.sock"
+NACK_SOCK_PATH = "unix:///tmp/async_map_stream_nack.sock"
+
+NACK_TEST_OPTIONS = NackOptions(delay=1000, max_deliveries=3, reason="retry")
 
 
 async def async_map_stream_handler(keys: list[str], datum: Datum) -> AsyncIterable[Message]:
@@ -85,7 +90,7 @@ def test_map_stream(map_stream_stub):
 
     # Prepare expected payload
     expected_payload = bytes(
-        "payload:test_mock_message " "event_time:2022-09-12 16:00:00 watermark:2022-09-12 16:01:00",
+        "payload:test_mock_message event_time:2022-09-12 16:00:00 watermark:2022-09-12 16:01:00",
         encoding="utf-8",
     )
 
@@ -106,16 +111,63 @@ def test_map_stream(map_stream_stub):
         result_msg_count += 1
 
     # Validate totals
-    assert (
-        result_msg_count == expected_result_msgs
-    ), f"Expected {expected_result_msgs} result messages, got {result_msg_count}"
+    assert result_msg_count == expected_result_msgs, (
+        f"Expected {expected_result_msgs} result messages, got {result_msg_count}"
+    )
     assert eot_count == expected_eots, f"Expected {expected_eots} EOT messages, got {eot_count}"
 
     # Validate 10 messages per request id: test-id-0..test-id-(req_count-1)
     for i in range(req_count):
-        assert (
-            id_counter[f"test-id-{i}"] == 10
-        ), f"Expected 10 results for test-id-{i}, got {id_counter[f'test-id-{i}']}"
+        assert id_counter[f"test-id-{i}"] == 10, (
+            f"Expected 10 results for test-id-{i}, got {id_counter[f'test-id-{i}']}"
+        )
+
+
+async def async_nack_map_stream_handler(keys: list[str], datum: Datum) -> AsyncIterable[Message]:
+    for i in range(3):
+        yield Message.to_nack(NACK_TEST_OPTIONS)
+
+
+async def _start_nack_server(udfs):
+    server = grpc.aio.server()
+    map_pb2_grpc.add_MapServicer_to_server(udfs, server)
+    server.add_insecure_port(NACK_SOCK_PATH)
+    logging.info("Starting nack server on %s", NACK_SOCK_PATH)
+    await server.start()
+    return server, NACK_SOCK_PATH
+
+
+@pytest.fixture(scope="module")
+def async_nack_map_stream_server():
+    loop = create_async_loop()
+    server_obj = MapStreamAsyncServer(map_stream_instance=async_nack_map_stream_handler)
+    udfs = server_obj.servicer
+    server = start_async_server(loop, _start_nack_server(udfs))
+    yield loop
+    teardown_async_server(loop, server)
+
+
+def test_map_stream_nack(async_nack_map_stream_server):
+    with grpc.insecure_channel(NACK_SOCK_PATH) as channel:
+        stub = map_pb2_grpc.MapStub(channel)
+        generator_response = stub.MapFn(request_iterator=request_generator(count=2, session=1))
+
+        handshake = next(generator_response)
+        assert handshake.handshake.sot
+
+        result_count = 0
+        for msg in generator_response:
+            if hasattr(msg, "status") and msg.status.eot:
+                continue
+            result = msg.results[0]
+            assert NACK in result.tags
+            assert result.nack_options.delay == NACK_TEST_OPTIONS.delay
+            assert result.nack_options.max_deliveries == NACK_TEST_OPTIONS.max_deliveries
+            assert result.nack_options.reason == NACK_TEST_OPTIONS.reason
+            result_count += 1
+
+    # 2 requests x 3 nack messages each
+    assert result_count == 6
 
 
 def test_is_ready(async_map_stream_server):