test(pyamber): cover the source and batch operator final methods

amber/src/test/python/core/models/test_operator.py

@@ -17,6 +17,7 @@
 
 import base64
 
+import pandas
 import pytest
 
 from core.models import (
@@ -51,6 +52,104 @@ def process_batch(self, batch, port):
         yield batch
 
 
+class _ProducingSource(SourceOperator):
+    """Source whose produce() actually emits raw (non-Tuple) records."""
+
+    def produce(self):
+        yield {"x": 1}
+        yield {"x": 2}
+
+
+class _TableProducingSource(SourceOperator):
+    """Source whose produce() emits a whole Table in one go."""
+
+    def produce(self):
+        yield Table([Tuple({"x": 1}), Tuple({"x": 2})])
+
+
+class _MixedProducingSource(SourceOperator):
+    """Source whose produce() interleaves a None signal between two records."""
+
+    def produce(self):
+        yield {"x": 1}
+        yield None
+        yield {"x": 2}
+
+
+class _NoneOutputBatch(BatchOperator):
+    """Batch operator whose process_batch declines to emit anything.
+
+    It records the Batches it is handed so a test asserting *absence* of output
+    can also assert the batch actually ran.
+    """
+
+    BATCH_SIZE = 1
+
+    def __init__(self):
+        super().__init__()
+        self.batches = []
+
+    def process_batch(self, batch, port):
+        self.batches.append(batch)
+        yield None
+
+
+class _TupleOutputBatch(BatchOperator):
+    """Batch operator whose process_batch emits a non-DataFrame output."""
+
+    BATCH_SIZE = 1
+
+    def process_batch(self, batch, port):
+        yield Tuple({"y": 42})
+
+
+class _DataFrameOutputBatch(BatchOperator):
+    """Batch operator whose process_batch emits a multi-row DataFrame."""
+
+    BATCH_SIZE = 1
+
+    def process_batch(self, batch, port):
+        yield pandas.DataFrame([{"y": 1}, {"y": 2}])
+
+
+class _MultiOutputBatch(BatchOperator):
+    """Batch operator whose process_batch emits more than one output."""
+
+    BATCH_SIZE = 1
+
+    def process_batch(self, batch, port):
+        yield Tuple({"y": 1})
+        yield Tuple({"y": 2})
+
+
+class _SpyBatch(BatchOperator):
+    """Batch operator that records the rows and the port of every Batch handed
+    to it, and echoes the row count back as its single output."""
+
+    BATCH_SIZE = 2
+
+    def __init__(self):
+        super().__init__()
+        self.seen = []
+        self.ports = []
+
+    def process_batch(self, batch, port):
+        self.seen.append([list(row) for _, row in batch.iterrows()])
+        self.ports.append(port)
+        yield Tuple({"batched": len(self.seen[-1])})
+
+
+def _take(iterator, limit):
+    """The first `limit` items of `iterator`.
+
+    Bounded on purpose: `BatchOperator.on_finish` loops `while` the port buffer
+    is non-empty, so a mutant that stops `_process_batch` from draining that
+    buffer turns it into an infinite generator. An unbounded `list()` would
+    hang there instead of failing an assertion.
+    """
+    return [item for _, item in zip(range(limit), iterator)]
+
+
 class _ConcreteTable(TableOperator):
     """Concrete subclass that records the table it received via process_table."""
 
@@ -179,6 +278,15 @@ def test_produce_state_on_start_returns_none_by_default(self):
     def test_produce_state_on_finish_returns_none_by_default(self):
         assert _ConcreteOperator().produce_state_on_finish(port=0) is None
 
+    def test_default_on_finish_yields_exactly_one_none(self):
+        # TupleOperatorV2.on_finish is the model-layer default hook. No in-repo
+        # subclass inherits it: SourceOperator/BatchOperator/TableOperator all
+        # override it, and pytexera's UDFOperatorV2 shadows it with a
+        # byte-identical body of its own. So this pins the base-class contract
+        # for any future direct subclass -- a generator that yields exactly one
+        # None, not an empty generator and not a plain return.
+        assert list(_ConcreteOperator().on_finish(port=0)) == [None]
+
 
 class TestLazyTemplateDecoder:
     def test_first_call_creates_decoder_and_caches_on_instance(self):
@@ -288,3 +396,121 @@ def test_buffers_are_keyed_by_port(self):
         list(op.on_finish(port=0))
         rows = list(op.received_tables[0].as_tuples())
         assert rows == [Tuple({"x": 1})]
+
+
+class TestSourceOperatorFinalMethods:
+    """SourceOperator replaces both TupleOperatorV2 tuple hooks: on_finish is a
+    source's only output path, and process_tuple is deliberately inert because
+    a source has no input."""
+
+    def test_on_finish_converts_each_produced_item_to_tuples(self):
+        # produce() may emit raw records; on_finish must normalize every one of
+        # them into a Tuple before it leaves the operator. Tuple.__eq__ starts
+        # with `isinstance(other, Tuple)`, so this equality already pins the
+        # type -- a raw dict fails it outright.
+        out = list(_ProducingSource().on_finish(port=0))
+        assert out == [Tuple({"x": 1}), Tuple({"x": 2})]
+
+    def test_on_finish_flattens_a_produced_table_into_its_tuples(self):
+        # A single produced Table must be exploded into its rows, not passed
+        # downstream as one Table object.
+        out = list(_TableProducingSource().on_finish(port=0))
+        assert out == [Tuple({"x": 1}), Tuple({"x": 2})]
+
+    def test_on_finish_preserves_produce_order_and_forwards_a_none_signal(self):
+        # produce() yielding None is the documented "no data this round" signal.
+        # Placing it *between* two records pins three things at once: the None
+        # survives conversion, the records around it are still normalized, and
+        # produce()'s emission order is preserved.
+        out = list(_MixedProducingSource().on_finish(port=0))
+        assert out == [Tuple({"x": 1}), None, Tuple({"x": 2})]
+
+    def test_process_tuple_yields_exactly_one_none(self):
+        # A source ignores any tuple handed to it, but still has to behave as a
+        # generator producing one None.
+        src = _ConcreteSource()
+        assert list(src.process_tuple(Tuple({"x": 1}), port=0)) == [None]
+
+
+class TestBatchOperatorOutputConversion:
+    """_process_batch owns the batch->output conversion: drop Nones, explode
+    DataFrames row-wise, pass anything else through untouched."""
+
+    def test_none_output_batch_is_dropped(self):
+        # BATCH_SIZE=1 makes process_tuple flush immediately.
+        op = _NoneOutputBatch()
+        assert list(op.process_tuple(Tuple({"x": 1}), port=0)) == []
+        # Positive control for the absence above: the empty output must mean
+        # "the None was dropped", not "no batch ever ran". A mutant that
+        # suppresses the flush, or never calls process_batch, leaves this empty.
+        assert len(op.batches) == 1
+        assert list(op.batches[0].columns) == ["x"]
+
+    def test_non_dataframe_output_batch_is_yielded_as_is(self):
+        op = _TupleOutputBatch()
+        out = list(op.process_tuple(Tuple({"x": 1}), port=0))
+        assert out == [Tuple({"y": 42})]
+
+    def test_dataframe_output_batch_is_exploded_into_rows(self):
+        # The counterweight to the test above: a DataFrame output must come out
+        # as one object per row, not as the frame itself.
+        op = _DataFrameOutputBatch()
+        out = list(op.process_tuple(Tuple({"x": 1}), port=0))
+        assert len(out) == 2
+        # This also pins provenance: the rows come from the *output* frame
+        # (column "y"), not the input batch (column "x"). A mutant that
+        # iterates the input batch fails here with a KeyError, or on the row
+        # count above. An extra `isinstance(row, DataFrame)` check would be
+        # dead weight -- DataFrame.iterrows() yields Series by library
+        # contract, and the frame-yielding mutant fails the count first.
+        assert [row["y"] for row in out] == [1, 2]
+
+    def test_every_output_of_process_batch_is_forwarded_in_order(self):
+        # One Batch may produce several outputs; all of them must come out, in
+        # the order process_batch emitted them.
+        op = _MultiOutputBatch()
+        out = list(op.process_tuple(Tuple({"x": 1}), port=0))
+        assert out == [Tuple({"y": 1}), Tuple({"y": 2})]
+
+
+class TestBatchOperatorBatchAssembly:
+    """The other half of _process_batch: how the Batch handed to process_batch
+    is assembled out of the per-port tuple buffer."""
+
+    def test_flush_hands_process_batch_the_buffered_rows_in_order(self):
+        op = _SpyBatch()  # BATCH_SIZE = 2
+        # First tuple is buffered only -- the batch is not full yet.
+        assert list(op.process_tuple(Tuple({"x": 1}), port=1)) == []
+        assert op.seen == []
+        # The second tuple fills the batch and triggers the flush.
+        out = list(op.process_tuple(Tuple({"x": 2}), port=1))
+        assert op.seen == [[[1], [2]]]  # FIFO: 1 before 2
+        assert op.ports == [1]  # the port it was buffered on, not port 0
+        assert out == [Tuple({"batched": 2})]
+
+    def test_on_finish_drains_the_buffer_in_batch_size_chunks(self):
+        op = _SpyBatch()
+        op.BATCH_SIZE = 5  # buffer more than one batch's worth without flushing
+        for value in (1, 2, 3):
+            assert list(op.process_tuple(Tuple({"x": value}), port=1)) == []
+        assert op.seen == []
+        op.BATCH_SIZE = 2  # now the buffer holds more than a single batch
+        out = _take(op.on_finish(port=1), 8)
+        # Two chunks: BATCH_SIZE rows, then the remainder. Not one big batch
+        # (the min() cap), not a single chunk (the while loop), and not LIFO.
+        assert op.seen == [[[1], [2]], [[3]]]
+        assert op.ports == [1, 1]
+        assert out == [Tuple({"batched": 2}), Tuple({"batched": 1})]
+
+    def test_batch_buffers_are_keyed_by_port(self):
+        # Mirrors TestTableOperator::test_buffers_are_keyed_by_port for the
+        # batch path: a tuple arriving on one port must not fill another
+        # port's batch.
+        op = _SpyBatch()  # BATCH_SIZE = 2
+        assert list(op.process_tuple(Tuple({"x": 1}), port=0)) == []
+        assert list(op.process_tuple(Tuple({"x": 99}), port=1)) == []
+        assert op.seen == []  # neither port reached 2 tuples
+        out = list(op.process_tuple(Tuple({"x": 2}), port=0))
+        assert op.seen == [[[1], [2]]]  # the port-1 tuple is not in here
+        assert op.ports == [0]
+        assert out == [Tuple({"batched": 2})]