diff --git a/amber/src/test/python/core/models/test_operator.py b/amber/src/test/python/core/models/test_operator.py index 99b3c729e85..bb18853350e 100644 --- a/amber/src/test/python/core/models/test_operator.py +++ b/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})]