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-6890-6b569dff00c0bbe2cb8e9c24a43d4aa20c56c528 in repository https://gitbox.apache.org/repos/asf/texera.git
commit cc6f09dedba417b1da4b0d59b52595f5d3d2da5d Author: Xinyuan Lin <[email protected]> AuthorDate: Sat Jul 25 19:18:27 2026 -0700 test(pyamber): extend network receiver and heartbeat runnable coverage (#6890) ### What changes were proposed in this PR? Extends the two sibling runnable suites to 39 tests, covering the uncovered dispatch, lifecycle and guard branches of `network_receiver.py` (~70% covered) and `heartbeat.py` (~62% covered). The behavior is driven deterministically — no real sleeps — and the tests leave no sockets or threads running. No source changes. ### Any related issues, documentation, discussions? Closes #6887. ### How was this PR tested? `python -m pytest src/test/python/core/runnables/test_network_receiver.py src/test/python/core/runnables/test_heartbeat.py` -> 39 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]) Co-authored-by: Meng Wang <[email protected]> --- .../test/python/core/runnables/test_heartbeat.py | 204 +++++++++++++- .../python/core/runnables/test_network_receiver.py | 295 ++++++++++++++++++++- 2 files changed, 496 insertions(+), 3 deletions(-) diff --git a/amber/src/test/python/core/runnables/test_heartbeat.py b/amber/src/test/python/core/runnables/test_heartbeat.py index 0fe5b321f5..76c3c3808b 100644 --- a/amber/src/test/python/core/runnables/test_heartbeat.py +++ b/amber/src/test/python/core/runnables/test_heartbeat.py @@ -17,7 +17,8 @@ import socket from threading import Event -from unittest.mock import patch, MagicMock +from types import SimpleNamespace +from unittest.mock import call, patch, MagicMock import pytest @@ -28,6 +29,36 @@ def make_heartbeat(host="localhost", port=12345, interval=0.05, event=None): return Heartbeat(host, port, interval, event or Event()) +class ScriptedEvent: + """A drop-in for ``threading.Event`` whose ``wait`` returns a scripted + sequence instead of sleeping. + + ``False`` means "the timeout elapsed, keep looping" and ``True`` means + "the stop event fired, leave the loop", which is exactly how + ``Heartbeat.run`` reads the return value. Driving the loop this way keeps + the tests deterministic and instantaneous -- no real sleeps -- while still + recording the timeout each call was made with, so the interval wiring + stays under test. Once the script is exhausted it returns ``True`` so a + regression can never spin forever. + """ + + def __init__(self, results): + self._results = list(results) + self.wait_timeouts = [] + + def wait(self, timeout=None): + self.wait_timeouts.append(timeout) + if not self._results: + return True + return self._results.pop(0) + + +# ``signal.SIGKILL`` does not exist on Windows, so `Heartbeat.stop` is pinned +# against these stand-ins: the assertions then check that the kill signals are +# read off the `signal` module rather than hard-coded, on every platform. +FAKE_SIGNALS = SimpleNamespace(SIGKILL="SIGKILL", SIGTERM="SIGTERM") + + class TestHeartbeatInit: def test_parses_host_and_port_from_grpc_tcp_url(self): hb = make_heartbeat(host="example.test", port=9090) @@ -111,3 +142,174 @@ class TestRunEarlyExit: def test_init_accepts_full_port_range(port): hb = make_heartbeat(port=port) assert hb._parsed_server_port == port + + +class TestRunLoop: + @pytest.mark.timeout(2) + def test_polls_once_per_interval_while_the_server_stays_up(self): + event = ScriptedEvent([False, False, True]) + hb = make_heartbeat(interval=0.25, event=event) + with ( + patch.object(hb, "_check_heartbeat", return_value=True) as mock_check, + patch.object(hb, "stop") as mock_stop, + ): + hb.run() + + # one probe per elapsed interval, and none after the stop event fired + assert mock_check.call_count == 2 + assert event.wait_timeouts == [0.25, 0.25, 0.25] + mock_stop.assert_not_called() + + @pytest.mark.timeout(2) + def test_a_single_failed_probe_is_not_enough_to_stop(self): + event = ScriptedEvent([False, False, True]) + hb = make_heartbeat(interval=0.05, event=event) + # first probe fails, the double check succeeds -> treated as a blip + with ( + patch.object( + hb, "_check_heartbeat", side_effect=[False, True, True] + ) as mock_check, + patch.object(hb, "stop") as mock_stop, + ): + hb.run() + + assert mock_check.call_count == 3 + mock_stop.assert_not_called() + + @pytest.mark.timeout(2) + def test_two_failed_probes_stop_the_worker_and_end_the_loop(self): + event = ScriptedEvent([False, False, False]) + hb = make_heartbeat(interval=0.05, event=event) + with ( + patch.object( + hb, "_check_heartbeat", side_effect=[False, False] + ) as mock_check, + patch.object(hb, "stop") as mock_stop, + ): + hb.run() + + assert mock_check.call_count == 2 + mock_stop.assert_called_once_with() + # run() returned instead of waiting again, even though the stop event + # was never set + assert event.wait_timeouts == [0.05] + + @pytest.mark.timeout(2) + def test_logs_the_unchanged_parent_pid_and_its_status_before_stopping(self): + with patch("core.runnables.heartbeat.os.getppid", return_value=4242): + hb = make_heartbeat(interval=0.05, event=ScriptedEvent([False])) + + with ( + patch("core.runnables.heartbeat.os.getppid", return_value=4242), + patch("core.runnables.heartbeat.psutil.Process") as mock_process, + patch("core.runnables.heartbeat.logger") as mock_logger, + patch.object(hb, "_check_heartbeat", return_value=False), + patch.object(hb, "stop"), + ): + mock_process.return_value.status.return_value = "zombie" + hb.run() + + mock_process.assert_called_once_with(4242) + message = mock_logger.warning.call_args[0][0] + assert "Parent process PID 4242 runs unusually." in message + assert "Parent PID hasn't changed." in message + assert "Original parent process Status: zombie" in message + + @pytest.mark.timeout(2) + def test_reports_a_reparented_and_vanished_parent_process(self): + with patch("core.runnables.heartbeat.os.getppid", return_value=4242): + hb = make_heartbeat(interval=0.05, event=ScriptedEvent([False])) + + with ( + # the JVM died and this worker was reparented to init + patch("core.runnables.heartbeat.os.getppid", return_value=1), + patch( + "core.runnables.heartbeat.psutil.Process", + side_effect=RuntimeError("no such process"), + ), + patch("core.runnables.heartbeat.logger") as mock_logger, + patch.object(hb, "_check_heartbeat", return_value=False), + patch.object(hb, "stop") as mock_stop, + ): + hb.run() + + message = mock_logger.warning.call_args[0][0] + assert "Parent PID changed to 1." in message + # a lookup failure degrades to a placeholder status instead of + # bubbling out of run() + assert "Original parent process Status: NOT FOUND" in message + mock_stop.assert_called_once_with() + + +class TestStop: + @staticmethod + def fake_child(pid, running=True): + child = MagicMock(name=f"child-{pid}") + child.pid = pid + child.is_running.return_value = running + return child + + def test_kills_running_descendants_then_terminates_itself(self): + hb = make_heartbeat() + running_child = self.fake_child(111) + dead_child = self.fake_child(222, running=False) + + with ( + patch("core.runnables.heartbeat.psutil.Process") as mock_process, + patch("core.runnables.heartbeat.os.kill") as mock_kill, + patch("core.runnables.heartbeat.os.getpid", return_value=777), + patch("core.runnables.heartbeat.signal", FAKE_SIGNALS), + ): + mock_process.return_value.children.return_value = [ + running_child, + dead_child, + ] + hb.stop() + + mock_process.return_value.children.assert_called_once_with(recursive=True) + # the already-dead child is skipped; self-termination comes last + assert mock_kill.call_args_list == [ + call(111, "SIGKILL"), + call(777, "SIGTERM"), + ] + + def test_terminates_itself_even_without_any_children(self): + hb = make_heartbeat() + with ( + patch("core.runnables.heartbeat.psutil.Process") as mock_process, + patch("core.runnables.heartbeat.os.kill") as mock_kill, + patch("core.runnables.heartbeat.os.getpid", return_value=777), + patch("core.runnables.heartbeat.signal", FAKE_SIGNALS), + ): + mock_process.return_value.children.return_value = [] + hb.stop() + + assert mock_kill.call_args_list == [call(777, "SIGTERM")] + + def test_a_failing_child_kill_does_not_abort_the_cleanup(self): + hb = make_heartbeat() + first_child = self.fake_child(111) + second_child = self.fake_child(222) + + with ( + patch("core.runnables.heartbeat.psutil.Process") as mock_process, + patch("core.runnables.heartbeat.os.kill") as mock_kill, + patch("core.runnables.heartbeat.os.getpid", return_value=777), + patch("core.runnables.heartbeat.signal", FAKE_SIGNALS), + patch("core.runnables.heartbeat.logger") as mock_logger, + ): + mock_process.return_value.children.return_value = [ + first_child, + second_child, + ] + mock_kill.side_effect = [ProcessLookupError("gone"), None, None] + hb.stop() + + assert mock_kill.call_args_list == [ + call(111, "SIGKILL"), + call(222, "SIGKILL"), + call(777, "SIGTERM"), + ] + warning = mock_logger.warning.call_args[0][0] + assert "PID 111" in warning + assert "gone" in warning diff --git a/amber/src/test/python/core/runnables/test_network_receiver.py b/amber/src/test/python/core/runnables/test_network_receiver.py index 49ba2408ef..8fbca258bb 100644 --- a/amber/src/test/python/core/runnables/test_network_receiver.py +++ b/amber/src/test/python/core/runnables/test_network_receiver.py @@ -17,8 +17,20 @@ import pytest import threading +from pampy import MatchError from pyarrow import Table +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from core.architecture.handlers.actorcommand.actor_handler_base import ( + ActorCommandHandler, +) +from core.architecture.handlers.actorcommand.backpressure_handler import ( + BackpressureHandler, +) +from core.architecture.handlers.actorcommand.credit_update_handler import ( + CreditUpdateHandler, +) from core.models.internal_queue import ( InternalQueue, DCMElement, @@ -27,7 +39,7 @@ from core.models.internal_queue import ( ) from core.models.payload import DataFrame, StateFrame from core.models.state import State -from core.proxy import ProxyClient +from core.proxy import ProxyClient, ProxyServer from core.runnables.network_receiver import NetworkReceiver from core.runnables.network_sender import NetworkSender from core.util.proto import set_one_of @@ -44,7 +56,15 @@ from proto.org.apache.texera.amber.engine.architecture.rpc import ( AsyncRpcContext, ControlRequest, ) -from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2 +from proto.org.apache.texera.amber.engine.common import ( + ActorCommand, + Backpressure, + CreditUpdate, + DirectControlMessagePayloadV2, + PythonActorMessage, + PythonControlMessage, + PythonDataHeader, +) class TestNetworkReceiver: @@ -239,3 +259,274 @@ class TestNetworkReceiver: assert element.payload.command_mapping == command_mapping assert element.payload.scope == scope assert element.tag == channel_id + + ################################################################### + # Socket-free tests. + # + # The fixtures below swap `ProxyServer` for a mock so the handlers + # registered by `NetworkReceiver.__init__` can be invoked directly with the + # exact bytes the Flight endpoints would hand them. That keeps the + # dispatch/lifecycle assertions hermetic (no port binding, no server + # thread) and lets us reach branches -- retry-on-bind-failure, unknown + # payload types, actor-command dispatch -- that the wire tests above + # cannot drive. + ################################################################### + + @pytest.fixture + def mock_server_class(self): + with patch("core.runnables.network_receiver.ProxyServer") as server_class: + # `register_shutdown` goes through `ProxyServer.ack`; keep the real + # decorator so the registered action's behavior stays under test. + server_class.ack = ProxyServer.ack + yield server_class + + @pytest.fixture + def offline_receiver(self, output_queue, mock_server_class): + """A NetworkReceiver whose ProxyServer is a mock, plus the three + handlers it registered on that server.""" + receiver = NetworkReceiver(output_queue, host="localhost", port=6666) + server = mock_server_class.return_value + return SimpleNamespace( + receiver=receiver, + server=server, + queue=output_queue, + data_handler=server.register_data_handler.call_args[0][0], + control_handler=server.register_control_handler.call_args[0][0], + actor_handler=server.register_actor_message_handler.call_args[0][0], + ) + + @staticmethod + def data_channel(): + worker_id = ActorVirtualIdentity(name="test") + return ChannelIdentity(worker_id, worker_id, False) + + @staticmethod + def control_channel(): + worker_id = ActorVirtualIdentity(name="test") + return ChannelIdentity(worker_id, worker_id, True) + + def test_init_retries_until_the_proxy_server_binds(self, output_queue): + with patch("core.runnables.network_receiver.ProxyServer") as server_class: + bound_server = MagicMock(name="bound_server") + server_class.side_effect = [OSError("port already in use"), bound_server] + receiver = NetworkReceiver(output_queue, host="localhost", port=6667) + + assert server_class.call_count == 2 + assert server_class.call_args_list[0].kwargs == { + "host": "localhost", + "port": 6667, + } + # The receiver keeps the instance from the successful attempt only. + assert receiver.proxy_server is bound_server + + def test_run_serves_on_the_proxy_server(self, offline_receiver): + offline_receiver.receiver.run() + offline_receiver.server.serve.assert_called_once_with() + + def test_stop_shuts_the_server_down_then_waits_for_it(self, offline_receiver): + offline_receiver.receiver.stop() + offline_receiver.server.graceful_shutdown.assert_called_once_with() + offline_receiver.server.wait.assert_called_once_with() + called = [name for name, _, _ in offline_receiver.server.mock_calls] + # waiting before shutting down would block forever + assert called.index("graceful_shutdown") < called.index("wait") + + def test_register_shutdown_registers_an_acking_shutdown_action( + self, offline_receiver + ): + shutdown = MagicMock(name="shutdown") + offline_receiver.receiver.register_shutdown(shutdown) + + offline_receiver.server.register.assert_called_once() + kwargs = offline_receiver.server.register.call_args.kwargs + assert kwargs["name"] == "shutdown" + shutdown.assert_not_called() + # The registered action must invoke the callback and ack the caller. + assert kwargs["action"]() == "Bye bye!" + shutdown.assert_called_once_with() + + @pytest.mark.timeout(10) + def test_data_handler_enqueues_a_data_frame_and_reports_credits( + self, offline_receiver, data_payload + ): + channel_id = self.data_channel() + header = PythonDataHeader(tag=channel_id, payload_type="Data") + + credits = offline_receiver.data_handler(bytes(header), data_payload.frame) + + assert credits == offline_receiver.queue.in_mem_size() + assert credits > 0 + element = offline_receiver.queue.get() + assert isinstance(element, DataElement) + assert isinstance(element.payload, DataFrame) + assert element.payload.frame is data_payload.frame + assert element.tag == channel_id + # is_control is normalized to a real bool so the tag hashes + # consistently as an internal-queue key. + assert element.tag.is_control is False + assert hash(element.tag) == hash(channel_id) + + @pytest.mark.timeout(10) + def test_data_handler_rebuilds_a_state_frame_from_its_columns( + self, offline_receiver + ): + channel_id = self.data_channel() + header = PythonDataHeader(tag=channel_id, payload_type="State") + table = Table.from_pydict( + { + State.CONTENT: [State({"answer": 42}).to_json()], + State.LOOP_COUNTER: [7], + State.LOOP_START_ID: ["loop-start-1"], + } + ) + + offline_receiver.data_handler(bytes(header), table) + + element = offline_receiver.queue.get() + assert isinstance(element.payload, StateFrame) + assert element.payload.frame == {"answer": 42} + assert element.payload.loop_counter == 7 + assert element.payload.loop_start_id == "loop-start-1" + assert element.tag == channel_id + + @pytest.mark.timeout(10) + def test_data_handler_wraps_an_ecm_payload_in_an_ecm_element( + self, offline_receiver + ): + channel_id = self.data_channel() + header = PythonDataHeader(tag=channel_id, payload_type="ECM") + ecm = EmbeddedControlMessage( + EmbeddedControlMessageIdentity("ecm-1"), + EmbeddedControlMessageType.NO_ALIGNMENT, + [self.control_channel()], + {}, + ) + table = Table.from_pydict({"payload": [bytes(ecm)]}) + + offline_receiver.data_handler(bytes(header), table) + + element = offline_receiver.queue.get() + assert isinstance(element, ECMElement) + assert element.payload == ecm + assert element.tag == channel_id + # every scope channel gets the same bool normalization as the tag + assert [c.is_control for c in element.payload.scope] == [True] + assert all(type(c.is_control) is bool for c in element.payload.scope) + + def test_data_handler_rejects_an_unknown_payload_type( + self, offline_receiver, data_payload + ): + header = PythonDataHeader(tag=self.data_channel(), payload_type="Nonsense") + + with pytest.raises(MatchError): + offline_receiver.data_handler(bytes(header), data_payload.frame) + + assert offline_receiver.queue.is_empty() + + @pytest.mark.timeout(10) + def test_control_handler_enqueues_a_dcm_element_without_charging_credits( + self, offline_receiver + ): + channel_id = self.control_channel() + payload = set_one_of(DirectControlMessagePayloadV2, ControlInvocation()) + message = PythonControlMessage(tag=channel_id, payload=payload) + + credits = offline_receiver.control_handler(bytes(message)) + + # in_mem_size only counts data channels, so a control message + # never consumes sender credits. + assert credits == 0 + element = offline_receiver.queue.get() + assert isinstance(element, DCMElement) + assert element.tag == channel_id + assert element.payload == payload + + def test_actor_handler_disables_data_when_backpressure_is_enabled( + self, offline_receiver, data_payload + ): + offline_receiver.queue.put( + DataElement(tag=self.data_channel(), payload=data_payload) + ) + assert offline_receiver.queue.is_data_enabled() + message = PythonActorMessage( + payload=set_one_of(ActorCommand, Backpressure(enable_backpressure=True)) + ) + + credits = offline_receiver.actor_handler(bytes(message)) + + assert credits == offline_receiver.queue.in_mem_size() + assert not offline_receiver.queue.is_data_enabled() + # backpressure only gates the data queues, it enqueues nothing + assert offline_receiver.queue.size_control() == 0 + + @pytest.mark.timeout(10) + def test_actor_handler_reenables_data_and_wakes_the_worker( + self, offline_receiver, data_payload + ): + offline_receiver.queue.put( + DataElement(tag=self.data_channel(), payload=data_payload) + ) + offline_receiver.queue.disable_data( + InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE + ) + message = PythonActorMessage( + payload=set_one_of(ActorCommand, Backpressure(enable_backpressure=False)) + ) + + offline_receiver.actor_handler(bytes(message)) + + assert offline_receiver.queue.is_data_enabled() + # a NoOperation control invocation is injected so the blocked main + # loop wakes up and drains the re-enabled data queues + assert offline_receiver.queue.size_control() == 1 + drained = [offline_receiver.queue.get() for _ in range(2)] + control_elements = [e for e in drained if isinstance(e, DCMElement)] + assert len(control_elements) == 1 + assert control_elements[0].tag.is_control + assert ( + control_elements[0].payload.control_invocation.method_name == "NoOperation" + ) + + def test_actor_handler_treats_a_credit_update_as_a_no_op(self, offline_receiver): + message = PythonActorMessage(payload=set_one_of(ActorCommand, CreditUpdate())) + + credits = offline_receiver.actor_handler(bytes(message)) + + assert credits == 0 + assert offline_receiver.queue.is_empty() + + def test_look_up_returns_the_handler_registered_for_the_command_type( + self, offline_receiver + ): + receiver = offline_receiver.receiver + assert isinstance(receiver.look_up(Backpressure()), BackpressureHandler) + assert isinstance(receiver.look_up(CreditUpdate()), CreditUpdateHandler) + + def test_look_up_raises_for_an_unregistered_command_type(self, offline_receiver): + with pytest.raises(KeyError): + offline_receiver.receiver.look_up(ActorCommand()) + + def test_register_actor_command_handler_keys_on_cmd_and_overwrites( + self, offline_receiver + ): + class RecordingCreditUpdateHandler(ActorCommandHandler): + cmd = CreditUpdate + + def __init__(self): + self.calls = [] + + def __call__(self, command, input_queue, *args, **kwargs): + self.calls.append(command) + + replacement = RecordingCreditUpdateHandler() + offline_receiver.receiver.register_actor_command_handler(replacement) + + assert offline_receiver.receiver.look_up(CreditUpdate()) is replacement + # the other registration is untouched + assert isinstance( + offline_receiver.receiver.look_up(Backpressure()), BackpressureHandler + ) + + message = PythonActorMessage(payload=set_one_of(ActorCommand, CreditUpdate())) + offline_receiver.actor_handler(bytes(message)) + assert replacement.calls == [CreditUpdate()]
