This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/texera.git
The following commit(s) were added to refs/heads/main by this push:
new ab58d477ba test(pyamber): extend core model, proxy and start-worker
handler coverage (#6988)
ab58d477ba is described below
commit ab58d477ba77d1f42cbe56f888dda96fe3fffad0
Author: Xinyuan Lin <[email protected]>
AuthorDate: Tue Jul 28 21:51:11 2026 -0700
test(pyamber): extend core model, proxy and start-worker handler coverage
(#6988)
### What changes were proposed in this PR?
Extends two pyamber suites and adds one new test file (59 tests total).
No source changes.
- **`core/models/single_blocking_io.py`** (~70% covered): extends the
existing suite over the uncovered branches.
- **`core/proxy/proxy_server.py`** (~85% covered): extends the existing
suite.
- **`core/architecture/handlers/control/start_worker_handler.py`** (~41%
covered, **no dedicated test**): new test file following the sibling
`test_end_worker_handler.py` pattern for constructing and asserting
against the handler's context/state manager.
Tests are hermetic and deterministic — no real sleeps, and no sockets or
threads left running.
### Any related issues, documentation, discussions?
Closes #6986.
### How was this PR tested?
`python -m pytest` over the three files -> 59 passed. `ruff format
--check` + `ruff check` clean.
### Was this PR authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8 [1M context])
---
.../handlers/control/test_start_worker_handler.py | 218 +++++++++++++++++++++
.../python/core/models/test_single_blocking_io.py | 114 +++++++++++
.../test/python/core/proxy/test_proxy_server.py | 185 ++++++++++++++++-
3 files changed, 515 insertions(+), 2 deletions(-)
diff --git
a/amber/src/test/python/core/architecture/handlers/control/test_start_worker_handler.py
b/amber/src/test/python/core/architecture/handlers/control/test_start_worker_handler.py
new file mode 100644
index 0000000000..4d1afb30ef
--- /dev/null
+++
b/amber/src/test/python/core/architecture/handlers/control/test_start_worker_handler.py
@@ -0,0 +1,218 @@
+# 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 asyncio
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from core.architecture.handlers.control.start_worker_handler import
StartWorkerHandler
+from core.architecture.managers.context import WORKER_STATE_TRANSITIONS
+from core.architecture.managers.state_manager import (
+ InvalidTransitionException,
+ StateManager,
+)
+from core.architecture.packaging.input_manager import InputManager
+from core.models.internal_queue import ECMElement, InternalQueue
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity,
+ PortIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ EmbeddedControlMessageType,
+ EmptyRequest,
+ WorkerStateResponse,
+)
+from proto.org.apache.texera.amber.engine.architecture.worker import
WorkerState
+
+WORKER_ID = "worker-1"
+
+SOURCE_CHANNEL = ChannelIdentity(
+ InputManager.SOURCE_STARTER, ActorVirtualIdentity(WORKER_ID), False
+)
+SOURCE_PORT = PortIdentity(0, False)
+
+
+def make_context(
+ is_source: bool,
+ initial_state: WorkerState = WorkerState.READY,
+ input_manager=None,
+):
+ """Wire the slice of Context that StartWorkerHandler touches.
+
+ The state manager and the input queue are the real ones so state
+ transitions and the self-fed source markers are asserted against real
+ behavior; only the executor is stubbed, since the handler reads nothing
+ from it but `is_source`.
+ """
+ input_queue = InternalQueue()
+ return SimpleNamespace(
+ worker_id=WORKER_ID,
+ input_queue=input_queue,
+
executor_manager=SimpleNamespace(executor=SimpleNamespace(is_source=is_source)),
+ input_manager=(
+ input_manager
+ if input_manager is not None
+ else InputManager(WORKER_ID, input_queue)
+ ),
+ state_manager=StateManager(WORKER_STATE_TRANSITIONS, initial_state),
+ current_input_channel_id=None,
+ )
+
+
+def start(context) -> WorkerStateResponse:
+ return
asyncio.run(StartWorkerHandler(context).start_worker(EmptyRequest()))
+
+
+def drain(input_queue: InternalQueue):
+ return [input_queue.get() for _ in range(input_queue.size())]
+
+
+class TestStartWorkerOnSource:
+ def test_transits_to_running_and_reports_the_bumped_version(self):
+ context = make_context(is_source=True)
+
+ response = start(context)
+
+ assert isinstance(response, WorkerStateResponse)
+ assert response.state == WorkerState.RUNNING
+ # A source worker really transitions here, so the state manager's
+ # logical clock must advance and be reported with the new state.
+ assert response.state_version == 1
+ assert context.state_manager.get_current_state() == WorkerState.RUNNING
+
+ def test_registers_the_source_starter_channel_on_port_zero(self):
+ context = make_context(is_source=True)
+
+ start(context)
+
+ # A source has no upstream, so the handler fabricates port 0 and a
+ # channel from the synthetic SOURCE_STARTER actor to itself.
+ assert context.current_input_channel_id == SOURCE_CHANNEL
+ assert context.input_manager.get_port_id(SOURCE_CHANNEL) == SOURCE_PORT
+ port = context.input_manager.get_port(SOURCE_PORT)
+ assert port.get_channels() == {SOURCE_CHANNEL}
+ # The synthetic port carries no attributes and no materialized readers.
+ assert port.get_schema().get_attr_names() == []
+ assert context.input_manager.get_input_port_mat_reader_threads() == {
+ SOURCE_PORT: []
+ }
+
+ def test_self_feeds_start_then_end_channel_markers(self):
+ context = make_context(is_source=True)
+
+ start(context)
+
+ elements = drain(context.input_queue)
+ assert len(elements) == 2
+ assert all(isinstance(element, ECMElement) for element in elements)
+ assert [element.tag for element in elements] == [SOURCE_CHANNEL] * 2
+
+ start_ecm, end_ecm = (element.payload for element in elements)
+ # The start marker opens the channel without alignment; the end marker
+ # must be port-aligned so downstream completion is ordered correctly.
+ assert start_ecm.id == EmbeddedControlMessageIdentity("StartChannel")
+ assert start_ecm.ecm_type == EmbeddedControlMessageType.NO_ALIGNMENT
+ assert end_ecm.id == EmbeddedControlMessageIdentity("EndChannel")
+ assert end_ecm.ecm_type == EmbeddedControlMessageType.PORT_ALIGNMENT
+
+ def test_markers_carry_a_command_addressed_to_this_worker(self):
+ context = make_context(is_source=True)
+
+ start(context)
+
+ start_ecm, end_ecm = (element.payload for element in
drain(context.input_queue))
+ for ecm, method in ((start_ecm, "StartChannel"), (end_ecm,
"EndChannel")):
+ assert ecm.scope == []
+ # The mapping is keyed by the receiving worker, which is this
+ # worker itself since the source feeds its own input queue.
+ assert list(ecm.command_mapping) == [WORKER_ID]
+ invocation = ecm.command_mapping[WORKER_ID]
+ assert invocation.method_name == method
+ assert invocation.command_id == -1
+
+ def test_a_started_source_is_left_running_on_a_second_start(self):
+ # transit_to() is a no-op when already in the target state, so a
+ # repeated start must not raise nor bump the reported version.
+ context = make_context(is_source=True)
+ start(context)
+
+ response = start(context)
+
+ assert response.state == WorkerState.RUNNING
+ assert response.state_version == 1
+
+ def test_rejects_a_start_before_the_worker_is_ready(self):
+ # UNINITIALIZED -> RUNNING is not a legal edge; the handler must fail
+ # loudly instead of half-starting the source.
+ context = make_context(is_source=True,
initial_state=WorkerState.UNINITIALIZED)
+
+ with pytest.raises(InvalidTransitionException):
+ start(context)
+
+ assert context.state_manager.get_current_state() ==
WorkerState.UNINITIALIZED
+ assert context.input_queue.size() == 0
+ assert context.current_input_channel_id is None
+
+
+class TestStartWorkerOnNonSource:
+ def test_starts_materialization_reader_threads_when_present(self):
+ input_manager = MagicMock(spec=InputManager)
+ input_manager.get_input_port_mat_reader_threads.return_value = {
+ SOURCE_PORT: [object()]
+ }
+ context = make_context(is_source=False, input_manager=input_manager)
+
+ response = start(context)
+
+
input_manager.start_input_port_mat_reader_threads.assert_called_once_with()
+ # Reading from a materialized port does not start the worker; the
+ # main loop transitions on the first arriving tuple instead.
+ assert response.state == WorkerState.READY
+ assert response.state_version == 0
+ assert context.input_queue.size() == 0
+ assert context.current_input_channel_id is None
+
+ def test_does_nothing_without_materialization_reader_threads(self):
+ input_manager = MagicMock(spec=InputManager)
+ input_manager.get_input_port_mat_reader_threads.return_value = {}
+ context = make_context(is_source=False, input_manager=input_manager)
+
+ response = start(context)
+
+ input_manager.start_input_port_mat_reader_threads.assert_not_called()
+ input_manager.add_input_port.assert_not_called()
+ input_manager.register_input.assert_not_called()
+ assert response.state == WorkerState.READY
+ assert response.state_version == 0
+ assert context.input_queue.size() == 0
+
+ def test_reports_a_paused_state_untouched(self):
+ # The handler is a pure reporter for a non-source worker: whatever
+ # state the worker is in is echoed back with its version.
+ input_manager = MagicMock(spec=InputManager)
+ input_manager.get_input_port_mat_reader_threads.return_value = {}
+ context = make_context(is_source=False, input_manager=input_manager)
+ context.state_manager.transit_to(WorkerState.PAUSED)
+
+ response = start(context)
+
+ assert response.state == WorkerState.PAUSED
+ assert response.state_version == 1
diff --git a/amber/src/test/python/core/models/test_single_blocking_io.py
b/amber/src/test/python/core/models/test_single_blocking_io.py
index 1e284b4198..8969a6f368 100644
--- a/amber/src/test/python/core/models/test_single_blocking_io.py
+++ b/amber/src/test/python/core/models/test_single_blocking_io.py
@@ -17,6 +17,8 @@
from threading import Condition
+import pytest
+
from core.models.single_blocking_io import SingleBlockingIO
@@ -114,3 +116,115 @@ class TestReadline:
assert cond.notify_calls >= 1
# The value is cleared after being handed out.
assert io.value is None
+
+ def test_keeps_waiting_across_spurious_wakeups(self):
+ # A wakeup that does not publish a value must not end the loop:
+ # readline() has to re-check `value` and park again, otherwise pdb
+ # would receive None as a command line.
+ cond = ScriptedCondition()
+ io = SingleBlockingIO(cond)
+
+ def produce_on_third_wakeup():
+ if cond.wait_calls == 3:
+ io.write("late")
+ io.flush()
+
+ cond.on_wait = produce_on_third_wakeup
+
+ assert io.readline() == "late\n"
+ assert cond.wait_calls == 3
+ # Every parking round also notifies the producer side.
+ assert cond.notify_calls == 3
+
+ def test_ignores_a_limit_argument(self):
+ # pdb calls readline() through the IO API, which may pass a limit;
+ # the value is always handed out whole regardless.
+ io = SingleBlockingIO(Condition())
+ io.write("abcdef")
+ io.flush()
+
+ assert io.readline(2) == "abcdef\n"
+
+ def test_clears_the_value_even_when_waiting_raises(self):
+ # The clear lives in a `finally`, so an interrupted wait must not
+ # leave a stale line behind for the next reader to pick up.
+ cond = ScriptedCondition()
+ io = SingleBlockingIO(cond)
+
+ def explode():
+ io.value = "half-delivered\n"
+ raise KeyboardInterrupt
+
+ cond.on_wait = explode
+
+ with pytest.raises(KeyboardInterrupt):
+ io.readline()
+
+ assert io.value is None
+
+ def test_consecutive_lines_are_delivered_in_order(self):
+ # The IO holds one line at a time; a write/flush/readline cycle must
+ # be repeatable without any state leaking between lines.
+ io = SingleBlockingIO(Condition())
+
+ for line in ("first", "second", "third"):
+ io.write(line)
+ io.flush()
+ assert io.readline() == f"{line}\n"
+ assert io.value is None
+ assert io.buf == ""
+
+ def test_a_flush_overwrites_an_unread_value(self):
+ # Documented single-element semantics: there is no queue, so a second
+ # flush before a read replaces the pending line.
+ io = SingleBlockingIO(Condition())
+ io.write("stale")
+ io.flush()
+ io.write("fresh")
+ io.flush()
+
+ assert io.readline() == "fresh\n"
+
+
+class TestUnsupportedIOOperations:
+ """The remaining IO methods are deliberate no-ops: pdb never calls them,
+ but SingleBlockingIO is handed to pdb as a full IO, so they must be inert
+ rather than raising or returning junk."""
+
+ @pytest.mark.parametrize(
+ "call",
+ [
+ pytest.param(lambda io: io.close(), id="close"),
+ pytest.param(lambda io: io.fileno(), id="fileno"),
+ pytest.param(lambda io: io.isatty(), id="isatty"),
+ pytest.param(lambda io: io.read(4), id="read"),
+ pytest.param(lambda io: io.readable(), id="readable"),
+ pytest.param(lambda io: io.readlines(4), id="readlines"),
+ pytest.param(lambda io: io.seek(0, 0), id="seek"),
+ pytest.param(lambda io: io.seekable(), id="seekable"),
+ pytest.param(lambda io: io.tell(), id="tell"),
+ pytest.param(lambda io: io.truncate(0), id="truncate"),
+ pytest.param(lambda io: io.writable(), id="writable"),
+ pytest.param(lambda io: io.writelines(["a", "b"]),
id="writelines"),
+ pytest.param(lambda io: io.__next__(), id="next"),
+ pytest.param(lambda io: io.__iter__(), id="iter"),
+ pytest.param(lambda io: io.__enter__(), id="enter"),
+ pytest.param(lambda io: io.__exit__(None, None, None), id="exit"),
+ ],
+ )
+ def test_no_op_methods_return_none_without_raising(self, call):
+ io = SingleBlockingIO(Condition())
+
+ assert call(io) is None
+
+ def test_no_op_methods_do_not_disturb_the_pending_line(self):
+ io = SingleBlockingIO(Condition())
+ io.write("payload")
+ io.flush()
+
+ io.close()
+ io.writelines(["ignored"])
+ io.truncate(0)
+
+ # None of the inert methods may consume or mutate the buffered line.
+ assert io.readline() == "payload\n"
diff --git a/amber/src/test/python/core/proxy/test_proxy_server.py
b/amber/src/test/python/core/proxy/test_proxy_server.py
index 3055ce0c64..59d1f43bde 100644
--- a/amber/src/test/python/core/proxy/test_proxy_server.py
+++ b/amber/src/test/python/core/proxy/test_proxy_server.py
@@ -16,12 +16,13 @@
# under the License.
import threading
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
import pytest
+from pyarrow import Table
from pyarrow.flight import Action
-from core.proxy.proxy_server import ProxyServer
+from core.proxy.proxy_server import ProxyServer, get_free_local_port
class TestProxyServer:
@@ -85,3 +86,183 @@ class TestProxyServer:
next(results)
assert shutdown_started.wait(timeout=5)
mock_shutdown.assert_called_once()
+
+ @staticmethod
+ def do_put_args(command: bytes = b"cmd", data: Table = None):
+ """Build the (descriptor, reader, writer) triple that Flight hands to
+ do_put. Only the three attributes do_put touches are stubbed."""
+ if data is None:
+ data = Table.from_pydict({"x": [1, 2]})
+ descriptor = MagicMock()
+ descriptor.command = command
+ reader = MagicMock()
+ reader.read_all.return_value = data
+ return descriptor, reader, MagicMock()
+
+ def test_heartbeat_is_registered_and_acks(self, server):
+ # The default liveness action the client polls before sending work.
+ assert "heartbeat" in server._procedures
+ result = next(server.do_action(None, Action("heartbeat", b"")))
+ assert result.body.to_pybytes() == b"ack"
+
+ def test_unknown_action_raises_key_error(self, server):
+ with pytest.raises(KeyError, match="Unknown action"):
+ next(server.do_action(None, Action("no-such-action", b"")))
+
+ def test_registering_the_same_name_overwrites_the_previous_action(self,
server):
+ server.register("hello", lambda: "first")
+ server.register("hello", lambda: "second")
+
+ assert len(server._procedures) == 5 # 4 defaults + "hello", not 6
+ result = next(server.do_action(None, Action("hello", b"")))
+ assert result.body.to_pybytes() == b"second"
+
+ def test_action_returning_bytes_is_passed_through_unencoded(self, server):
+ # bytes results are forwarded verbatim; everything else is str()-ed.
+ server.register("raw", lambda: b"\x00\x01raw")
+
+ result = next(server.do_action(None, Action("raw", b"")))
+ assert result.body.to_pybytes() == b"\x00\x01raw"
+
+ def test_action_receives_the_body_as_payload_when_non_empty(self, server):
+ seen = []
+ server.register("echo", lambda payload: seen.append(payload) or
payload)
+
+ result = next(server.do_action(None, Action("echo", b"payload")))
+ assert seen == [b"payload"]
+ assert result.body.to_pybytes() == b"payload"
+
+ def test_error_inside_an_action_is_reraised_to_the_caller(self, server):
+ # register() wraps the action in a logging catcher, but with
+ # reraise=True so the failure still surfaces on the client call.
+ def boom():
+ raise ValueError("action failed")
+
+ server.register("boom", boom)
+ with pytest.raises(ValueError, match="action failed"):
+ next(server.do_action(None, Action("boom", b"")))
+
+ def test_list_actions_reports_names_with_descriptions(self, server):
+ server.register("hello", lambda: None, description="Say hello.")
+
+ listed = dict(server.list_actions(None))
+ assert listed["hello"] == "Say hello."
+ assert listed["shutdown"] == "Shut down this server."
+ assert listed["control"] == "Process the control message"
+ assert listed["actor"] == "Process the actor message"
+ # heartbeat is registered without a description.
+ assert listed["heartbeat"] == ""
+
+
+class TestProxyServerHandlers:
+ """Handler registration and the data path, which the default constructor
+ leaves unimplemented until the runtime wires them up."""
+
+ @pytest.fixture()
+ def server(self):
+ server = ProxyServer()
+ yield server
+ server.graceful_shutdown()
+
+ def test_data_handler_must_accept_command_and_data(self, server):
+ with pytest.raises(AssertionError):
+ server.register_data_handler(lambda only_one: None)
+
+ def test_control_handler_must_accept_at_least_one_argument(self, server):
+ with pytest.raises(AssertionError):
+ server.register_control_handler(lambda: None)
+
+ def test_do_put_without_a_registered_handler_is_not_implemented(self,
server):
+ descriptor, reader, writer = TestProxyServer.do_put_args()
+ with pytest.raises(NotImplementedError):
+ server.do_put(None, descriptor, reader, writer)
+
+ def test_control_action_without_a_registered_handler_is_not_implemented(
+ self, server
+ ):
+ with pytest.raises(NotImplementedError):
+ next(server.do_action(None, Action("control", b"payload")))
+
+ def
test_actor_action_without_a_registered_handler_is_not_implemented(self, server):
+ with pytest.raises(NotImplementedError):
+ next(server.do_action(None, Action("actor", b"payload")))
+
+ def test_do_put_routes_the_batch_and_writes_back_sender_credits(self,
server):
+ seen = []
+
+ def handler(command, data):
+ seen.append((command, data))
+ return 7
+
+ server.register_data_handler(handler)
+ data = Table.from_pydict({"x": [1, 2, 3]})
+ descriptor, reader, writer = TestProxyServer.do_put_args(b"cmd", data)
+
+ server.do_put(None, descriptor, reader, writer)
+
+ assert seen == [(b"cmd", data)]
+ # The credit is echoed back as a little-endian 8-byte integer.
+ written = writer.write.call_args[0][0]
+ assert written.to_pybytes() == (7).to_bytes(length=8,
byteorder="little")
+
+ def test_do_put_skips_the_credit_reply_for_a_non_int_result(self, server):
+ server.register_data_handler(lambda command, data: None)
+ descriptor, reader, writer = TestProxyServer.do_put_args()
+
+ server.do_put(None, descriptor, reader, writer)
+
+ writer.write.assert_not_called()
+
+ def test_control_action_is_routed_to_the_registered_control_handler(self,
server):
+ seen = []
+
+ def handler(control_message):
+ seen.append(control_message)
+ return 3
+
+ server.register_control_handler(handler)
+
+ result = next(server.do_action(None, Action("control",
b"control-bytes")))
+
+ assert seen == [b"control-bytes"]
+ # The reply carries the queue size used for credit calculation.
+ assert result.body.to_pybytes() == b"3"
+
+ def test_actor_action_is_routed_to_the_registered_actor_handler(self,
server):
+ seen = []
+
+ def handler(message):
+ seen.append(message)
+ return b"actor-ack"
+
+ server.register_actor_message_handler(handler)
+
+ result = next(server.do_action(None, Action("actor", b"actor-bytes")))
+
+ assert seen == [b"actor-bytes"]
+ assert result.body.to_pybytes() == b"actor-ack"
+
+
+class TestProxyServerPort:
+ def test_serves_on_the_requested_port(self):
+ port = get_free_local_port()
+ server = ProxyServer(port=port)
+ try:
+ assert server.get_port_number() == port
+ finally:
+ server.graceful_shutdown()
+
+ def test_picks_a_free_port_when_none_is_given(self):
+ server = ProxyServer()
+ try:
+ port = server.get_port_number()
+ assert isinstance(port, int)
+ assert 0 < port < 65536
+ # A second server must not collide with the first one's port.
+ other = ProxyServer()
+ try:
+ assert other.get_port_number() != port
+ finally:
+ other.graceful_shutdown()
+ finally:
+ server.graceful_shutdown()
