You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@texera.apache.org by gi...@apache.org on 2026/07/24 21:11:27 UTC
(texera) branch gh-readonly-queue/main/pr-6491-eae4f218a67978361496075da4cdbea07ee6cd74 created (now 28995b36f8)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6491-eae4f218a67978361496075da4cdbea07ee6cd74
in repository https://gitbox.apache.org/repos/asf/texera.git
at 28995b36f8 test(pyamber): add unit test coverage for InputManager (#6491)
This branch includes the following new commits:
new 28995b36f8 test(pyamber): add unit test coverage for InputManager (#6491)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(texera) 01/01: test(pyamber): add unit test coverage for InputManager (#6491)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6491-eae4f218a67978361496075da4cdbea07ee6cd74
in repository https://gitbox.apache.org/repos/asf/texera.git
commit 28995b36f8f5312efee4f957b1540a4a50a282da
Author: Eugene Gu <eu...@outlook.com>
AuthorDate: Fri Jul 24 14:11:05 2026 -0700
test(pyamber): add unit test coverage for InputManager (#6491)
### What changes were proposed in this PR?
Adds `test_input_manager.py` (new, 20 tests) for the Python worker's
`InputManager`
(`amber/src/main/python/core/architecture/packaging/input_manager.py`),
which had no direct unit tests — the Scala `InputManager` is covered by
#5451, but the Python counterpart was only exercised indirectly through
`MainLoop`. Uses a real `InternalQueue`, in-memory `Schema`/`pyarrow`
tables, and stubbed reader runnables, so no engine or storage is needed.
The suite covers:
- **Registration** — `add_input_port` / `register_input`: `None`
`id`/`internal` fields are normalized to `0`/`False`, re-adding an
existing port is ignored, and channel-to-port bindings are checked in
both directions.
- **Completion tracking** — completing a channel marks its port;
`all_ports_completed` is vacuously `True` with zero ports; completing
one channel completes the whole port even if its other channels are
still open (safe today because the `EndChannel` ECM is port-aligned —
documented in a test comment).
- **Channel filtering** — `get_all_data_channel_ids` excludes control
channels.
- **Materialization readers** — mismatched `uris`/`partitionings` raise;
one reader is created per URI with the right pairing; a second setup
replaces the readers instead of appending; finished readers are not
restarted, and unfinished ones run on daemon threads.
- **`process_data_payload`** — a DataFrame becomes `Tuple`s carrying the
port's schema; an empty table yields nothing; a StateFrame passes
through unchanged; a DataFrame from an unregistered channel fails only
when the generator is consumed; unknown payload types raise
`NotImplementedError`.
### Any related issues, documentation, discussions?
Closes #6486
### How was this PR tested?
```
cd amber/src/main/python
PYTHONPATH=. python -m pytest ../../test/python/core/architecture/packaging/ -v
# 28 passed (20 new InputManager tests + 8 existing OutputManager tests)
```
### Was this PR authored or co-authored using generative AI tooling?
Co-authored by: Claude Code (Fable 5)
---
.../architecture/packaging/test_input_manager.py | 409 +++++++++++++++++++++
1 file changed, 409 insertions(+)
diff --git a/amber/src/test/python/core/architecture/packaging/test_input_manager.py b/amber/src/test/python/core/architecture/packaging/test_input_manager.py
new file mode 100644
index 0000000000..6ba2a2c8f3
--- /dev/null
+++ b/amber/src/test/python/core/architecture/packaging/test_input_manager.py
@@ -0,0 +1,409 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import threading
+
+import pyarrow
+import pytest
+
+import core.architecture.packaging.input_manager as input_manager_module
+from core.architecture.packaging.input_manager import InputManager
+from core.models import InternalQueue, Schema
+from core.models.payload import DataFrame, DataPayload, StateFrame
+from core.models.state import State
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ PortIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ Partitioning,
+)
+
+WORKER_ID = "Worker:WF1-test-main-0"
+
+
+def _channel_id(from_name: str, is_control: bool = False) -> ChannelIdentity:
+ return ChannelIdentity(
+ ActorVirtualIdentity(from_name),
+ ActorVirtualIdentity(WORKER_ID),
+ is_control,
+ )
+
+
+class _RecordingReaderRunnable:
+ """Stands in for InputPortMaterializationReaderRunnable, whose real
+ constructor parses the URI and instantiates a partitioner. Records the
+ constructor arguments for assertions."""
+
+ def __init__(self, uri, queue, worker_actor_id, partitioning):
+ self.uri = uri
+ self.queue = queue
+ self.worker_actor_id = worker_actor_id
+ self.partitioning = partitioning
+
+
+class _StubReaderRunnable:
+ """A minimal runnable exposing just what start_input_port_mat_reader_threads
+ touches: finished(), run(), and channel_id (used in the thread name)."""
+
+ def __init__(self, finished: bool):
+ self._finished = finished
+ self.channel_id = "stub-channel"
+ self.ran = threading.Event()
+ self.ran_on_daemon_thread = None
+
+ def finished(self) -> bool:
+ return self._finished
+
+ def run(self) -> None:
+ self.ran_on_daemon_thread = threading.current_thread().daemon
+ self.ran.set()
+
+
+@pytest.fixture
+def input_queue():
+ return InternalQueue()
+
+
+@pytest.fixture
+def input_manager(input_queue):
+ return InputManager(WORKER_ID, input_queue)
+
+
+@pytest.fixture
+def schema():
+ return Schema(raw_schema={"name": "STRING", "age": "INTEGER"})
+
+
+class TestAddInputPort:
+ def test_none_port_fields_are_normalized(self, input_manager, schema):
+ port_id = PortIdentity(id=None, internal=None)
+
+ input_manager.add_input_port(port_id, schema, [], [])
+
+ # The passed-in identity is normalized in place, and the port is
+ # stored under the normalized key.
+ assert port_id.id == 0
+ assert port_id.internal is False
+ normalized = PortIdentity(id=0, internal=False)
+ assert input_manager.get_port(normalized).get_schema() is schema
+
+ def test_re_adding_same_port_keeps_original_worker_port(
+ self, input_manager, schema
+ ):
+ port_id = PortIdentity(id=0, internal=False)
+ input_manager.add_input_port(port_id, schema, [], [])
+ original_port = input_manager.get_port(port_id)
+
+ other_schema = Schema(raw_schema={"other": "STRING"})
+ input_manager.add_input_port(
+ PortIdentity(id=0, internal=False), other_schema, [], []
+ )
+
+ assert input_manager.get_port(port_id) is original_port
+ assert input_manager.get_port(port_id).get_schema() is schema
+
+ def test_none_normalized_port_dedups_against_existing_port(
+ self, input_manager, schema, monkeypatch
+ ):
+ monkeypatch.setattr(
+ input_manager_module,
+ "InputPortMaterializationReaderRunnable",
+ _RecordingReaderRunnable,
+ )
+ input_manager.add_input_port(
+ PortIdentity(id=0, internal=False), schema, ["uri-a"], [Partitioning()]
+ )
+ original_port = input_manager.get_port(PortIdentity(id=0, internal=False))
+
+ other_schema = Schema(raw_schema={"other": "STRING"})
+ input_manager.add_input_port(
+ PortIdentity(id=None, internal=None),
+ other_schema,
+ ["uri-b"],
+ [Partitioning()],
+ )
+
+ # The None-field identity normalizes to port 0, so this counts as a
+ # re-add: the original WorkerPort and its schema are kept ...
+ port_id = PortIdentity(id=0, internal=False)
+ assert input_manager.get_port(port_id) is original_port
+ assert input_manager.get_port(port_id).get_schema() is schema
+ # ... but set_up_input_port_mat_reader_threads is called
+ # unconditionally, so the port's reader list is still replaced.
+ readers = input_manager.get_input_port_mat_reader_threads()[port_id]
+ assert [reader.uri for reader in readers] == ["uri-b"]
+
+
+class TestRegisterInput:
+ def test_binds_channel_to_port_in_both_directions(self, input_manager, schema):
+ port_id = PortIdentity(id=1, internal=False)
+ channel_id = _channel_id("upstream-worker")
+ input_manager.add_input_port(port_id, schema, [], [])
+
+ input_manager.register_input(channel_id, port_id)
+
+ assert input_manager.get_port_id(channel_id) == port_id
+ assert channel_id in input_manager.get_port(port_id).get_channels()
+
+ def test_none_port_fields_are_normalized(self, input_manager, schema):
+ normalized = PortIdentity(id=0, internal=False)
+ channel_id = _channel_id("upstream-worker")
+ input_manager.add_input_port(PortIdentity(id=0, internal=False), schema, [], [])
+
+ input_manager.register_input(channel_id, PortIdentity(id=None, internal=None))
+
+ assert input_manager.get_port_id(channel_id) == normalized
+ assert channel_id in input_manager.get_port(normalized).get_channels()
+
+ def test_re_registering_channel_leaves_stale_reverse_mapping(
+ self, input_manager, schema
+ ):
+ port_a = PortIdentity(id=0, internal=False)
+ port_b = PortIdentity(id=1, internal=False)
+ channel_id = _channel_id("upstream-worker")
+ input_manager.add_input_port(port_a, schema, [], [])
+ input_manager.add_input_port(port_b, schema, [], [])
+
+ input_manager.register_input(channel_id, port_a)
+ input_manager.register_input(channel_id, port_b)
+
+ # The forward mapping is updated to the new port ...
+ assert input_manager.get_port_id(channel_id) == port_b
+ assert channel_id in input_manager.get_port(port_b).get_channels()
+ # ... but the old port's channel set is never cleaned up, so the
+ # channel remains in both reverse mappings (current behavior).
+ assert channel_id in input_manager.get_port(port_a).get_channels()
+
+
+class TestPortCompletion:
+ def test_completing_a_channel_marks_only_its_port(self, input_manager, schema):
+ port_a = PortIdentity(id=0, internal=False)
+ port_b = PortIdentity(id=1, internal=False)
+ channel_a = _channel_id("upstream-a")
+ channel_b = _channel_id("upstream-b")
+ input_manager.add_input_port(port_a, schema, [], [])
+ input_manager.add_input_port(port_b, schema, [], [])
+ input_manager.register_input(channel_a, port_a)
+ input_manager.register_input(channel_b, port_b)
+
+ input_manager.complete_current_port(channel_a)
+
+ assert input_manager.get_port(port_a).completed is True
+ assert input_manager.get_port(port_b).completed is False
+ assert input_manager.all_ports_completed() is False
+
+ input_manager.complete_current_port(channel_b)
+
+ assert input_manager.all_ports_completed() is True
+
+ def test_all_ports_completed_is_vacuously_true_without_ports(self, input_manager):
+ assert input_manager.all_ports_completed() is True
+
+ def test_completing_one_channel_completes_the_whole_port(
+ self, input_manager, schema
+ ):
+ # Completing ONE channel marks the WHOLE port as completed, even if
+ # the port's other channels are still open. InputManager itself never
+ # counts channels; it relies on its only caller, end_channel_handler,
+ # which runs once per port: the EndChannel ECM uses PORT_ALIGNMENT,
+ # so the handler fires only after every channel of the port has
+ # delivered the marker. If that alignment ever changes, completing a
+ # port this early would silently drop the open channels' data.
+ port_id = PortIdentity(id=0, internal=False)
+ channel_a = _channel_id("upstream-a")
+ channel_b = _channel_id("upstream-b")
+ input_manager.add_input_port(port_id, schema, [], [])
+ input_manager.register_input(channel_a, port_id)
+ input_manager.register_input(channel_b, port_id)
+
+ input_manager.complete_current_port(channel_a)
+
+ assert input_manager.get_port(port_id).completed is True
+ assert input_manager.all_ports_completed() is True
+
+
+class TestChannelIds:
+ def test_data_channel_ids_exclude_control_channels(self, input_manager, schema):
+ port_id = PortIdentity(id=0, internal=False)
+ data_channel = _channel_id("upstream-worker", is_control=False)
+ control_channel = _channel_id("controller", is_control=True)
+ input_manager.add_input_port(port_id, schema, [], [])
+ input_manager.register_input(data_channel, port_id)
+ input_manager.register_input(control_channel, port_id)
+
+ assert input_manager.get_all_data_channel_ids() == {data_channel}
+ assert set(input_manager.get_all_channel_ids()) == {
+ data_channel,
+ control_channel,
+ }
+
+
+class TestSetUpInputPortMatReaderThreads:
+ def test_mismatched_uris_and_partitionings_raise(self, input_manager):
+ with pytest.raises(AssertionError):
+ input_manager.set_up_input_port_mat_reader_threads(
+ PortIdentity(id=0, internal=False), ["uri-a"], []
+ )
+
+ def test_creates_one_reader_per_uri(self, input_manager, input_queue, monkeypatch):
+ monkeypatch.setattr(
+ input_manager_module,
+ "InputPortMaterializationReaderRunnable",
+ _RecordingReaderRunnable,
+ )
+ port_id = PortIdentity(id=0, internal=False)
+ partitionings = [Partitioning(), Partitioning()]
+
+ input_manager.set_up_input_port_mat_reader_threads(
+ port_id, ["uri-a", "uri-b"], partitionings
+ )
+
+ readers = input_manager.get_input_port_mat_reader_threads()[port_id]
+ assert [reader.uri for reader in readers] == ["uri-a", "uri-b"]
+ assert readers[0].partitioning is partitionings[0]
+ assert readers[1].partitioning is partitionings[1]
+ for reader in readers:
+ assert reader.queue is input_queue
+ assert reader.worker_actor_id == ActorVirtualIdentity(WORKER_ID)
+
+ def test_second_set_up_replaces_previous_readers(self, input_manager, monkeypatch):
+ monkeypatch.setattr(
+ input_manager_module,
+ "InputPortMaterializationReaderRunnable",
+ _RecordingReaderRunnable,
+ )
+ port_id = PortIdentity(id=0, internal=False)
+
+ input_manager.set_up_input_port_mat_reader_threads(
+ port_id, ["uri-a"], [Partitioning()]
+ )
+ input_manager.set_up_input_port_mat_reader_threads(
+ port_id, ["uri-b"], [Partitioning()]
+ )
+
+ # The port's reader list is replaced wholesale, not appended to.
+ readers = input_manager.get_input_port_mat_reader_threads()[port_id]
+ assert [reader.uri for reader in readers] == ["uri-b"]
+
+
+class TestStartInputPortMatReaderThreads:
+ def test_only_unfinished_readers_are_started(self, input_manager):
+ finished_reader = _StubReaderRunnable(finished=True)
+ unfinished_reader = _StubReaderRunnable(finished=False)
+ input_manager._input_port_mat_reader_runnables[
+ PortIdentity(id=0, internal=False)
+ ] = [finished_reader, unfinished_reader]
+
+ input_manager.start_input_port_mat_reader_threads()
+
+ assert unfinished_reader.ran.wait(timeout=5)
+ # The reader must run off the caller's thread, on a daemon thread.
+ assert unfinished_reader.ran_on_daemon_thread is True
+ assert not finished_reader.ran.is_set()
+
+
+class TestProcessDataPayload:
+ @pytest.fixture
+ def registered_channel(self, input_manager, schema):
+ port_id = PortIdentity(id=0, internal=False)
+ channel_id = _channel_id("upstream-worker")
+ input_manager.add_input_port(port_id, schema, [], [])
+ input_manager.register_input(channel_id, port_id)
+ return channel_id
+
+ def test_data_frame_yields_tuples_with_port_schema(
+ self, input_manager, schema, registered_channel
+ ):
+ table = pyarrow.Table.from_pydict(
+ {"name": ["Alice", "Bob"], "age": [10, 20]},
+ schema=schema.as_arrow_schema(),
+ )
+
+ results = list(
+ input_manager.process_data_payload(
+ registered_channel, DataFrame(frame=table)
+ )
+ )
+
+ assert [tuple_["name"] for tuple_ in results] == ["Alice", "Bob"]
+ assert [tuple_["age"] for tuple_ in results] == [10, 20]
+ # Each tuple is built against the schema of the channel's port.
+ for tuple_ in results:
+ assert tuple_._schema is schema
+
+ def test_state_frame_is_passed_through(self, input_manager, registered_channel):
+ state = State({"i": 2})
+
+ results = list(
+ input_manager.process_data_payload(
+ registered_channel, StateFrame(frame=state)
+ )
+ )
+
+ assert len(results) == 1
+ assert results[0] is state
+
+ def test_unknown_payload_type_raises(self, input_manager, registered_channel):
+ with pytest.raises(NotImplementedError):
+ list(input_manager.process_data_payload(registered_channel, DataPayload()))
+
+ def test_empty_data_frame_yields_no_tuples(
+ self, input_manager, schema, registered_channel
+ ):
+ empty_table = schema.as_arrow_schema().empty_table()
+
+ results = list(
+ input_manager.process_data_payload(
+ registered_channel, DataFrame(frame=empty_table)
+ )
+ )
+
+ assert results == []
+
+ def test_state_frame_through_unregistered_channel_passes_through(
+ self, input_manager
+ ):
+ # The StateFrame branch never touches the channel/port tables, so a
+ # channel that was never registered still passes state through.
+ state = State({"i": 2})
+
+ results = list(
+ input_manager.process_data_payload(
+ _channel_id("never-registered"), StateFrame(frame=state)
+ )
+ )
+
+ assert len(results) == 1
+ assert results[0] is state
+
+ def test_data_frame_through_unregistered_channel_fails_lazily(
+ self, input_manager, schema
+ ):
+ table = pyarrow.Table.from_pydict(
+ {"name": ["Alice"], "age": [10]},
+ schema=schema.as_arrow_schema(),
+ )
+
+ # process_data_payload is a generator: creating it raises nothing ...
+ generator = input_manager.process_data_payload(
+ _channel_id("never-registered"), DataFrame(frame=table)
+ )
+ # ... the unknown-channel lookup fails only upon iteration.
+ with pytest.raises(KeyError):
+ list(generator)