From e5949a727334f9301f42476b5edc6de58de36d95 Mon Sep 17 00:00:00 2001 From: Kasper Marinus Date: Tue, 25 Nov 2025 11:21:25 +0100 Subject: [PATCH] fix: fix test race condition ref: N25B-301 --- .../agents/actuation/robot_speech_agent.py | 2 +- .../bdi/bdi_core_agent/bdi_core_agent.py | 2 +- .../communication/ri_communication_agent.py | 2 +- .../transcription_agent.py | 2 +- .../agents/perception/vad_agent.py | 2 +- src/control_backend/core/agent_system.py | 8 +++--- .../actuation/test_robot_speech_agent.py | 27 +++---------------- .../test_ri_communication_agent.py | 26 +++--------------- test/unit/core/test_agent_system.py | 12 +++++---- 9 files changed, 25 insertions(+), 58 deletions(-) diff --git a/src/control_backend/agents/actuation/robot_speech_agent.py b/src/control_backend/agents/actuation/robot_speech_agent.py index 48316b9..e44f4bd 100644 --- a/src/control_backend/agents/actuation/robot_speech_agent.py +++ b/src/control_backend/agents/actuation/robot_speech_agent.py @@ -45,7 +45,7 @@ class RobotSpeechAgent(BaseAgent): self.subsocket.connect(settings.zmq_settings.internal_sub_address) self.subsocket.setsockopt(zmq.SUBSCRIBE, b"command") - await self.add_behavior(self._zmq_command_loop()) + self.add_behavior(self._zmq_command_loop()) self.logger.info("Finished setting up %s", self.name) diff --git a/src/control_backend/agents/bdi/bdi_core_agent/bdi_core_agent.py b/src/control_backend/agents/bdi/bdi_core_agent/bdi_core_agent.py index 72d3341..b798982 100644 --- a/src/control_backend/agents/bdi/bdi_core_agent/bdi_core_agent.py +++ b/src/control_backend/agents/bdi/bdi_core_agent/bdi_core_agent.py @@ -34,7 +34,7 @@ class BDICoreAgent(BaseAgent): await self._load_asl() # Start the BDI cycle loop - await self.add_behavior(self._bdi_loop()) + self.add_behavior(self._bdi_loop()) self._wake_bdi_loop.set() self.logger.debug("Setup complete.") diff --git a/src/control_backend/agents/communication/ri_communication_agent.py b/src/control_backend/agents/communication/ri_communication_agent.py index 8dfe368..50ea284 100644 --- a/src/control_backend/agents/communication/ri_communication_agent.py +++ b/src/control_backend/agents/communication/ri_communication_agent.py @@ -37,7 +37,7 @@ class RICommunicationAgent(BaseAgent): if await self._negotiate_connection(): self.connected = True - await self.add_behavior(self._listen_loop()) + self.add_behavior(self._listen_loop()) else: self.logger.warning("Failed to negotiate connection during setup.") diff --git a/src/control_backend/agents/perception/transcription_agent/transcription_agent.py b/src/control_backend/agents/perception/transcription_agent/transcription_agent.py index d3114ed..d0b0396 100644 --- a/src/control_backend/agents/perception/transcription_agent/transcription_agent.py +++ b/src/control_backend/agents/perception/transcription_agent/transcription_agent.py @@ -37,7 +37,7 @@ class TranscriptionAgent(BaseAgent): self.speech_recognizer.load_model() # Warmup # Start background loop - await self.add_behavior(self._transcribing_loop()) + self.add_behavior(self._transcribing_loop()) self.logger.info("Finished setting up %s", self.name) diff --git a/src/control_backend/agents/perception/vad_agent.py b/src/control_backend/agents/perception/vad_agent.py index ab6d6c7..374ffa6 100644 --- a/src/control_backend/agents/perception/vad_agent.py +++ b/src/control_backend/agents/perception/vad_agent.py @@ -93,7 +93,7 @@ class VADAgent(BaseAgent): # Warmup/reset await self.reset_stream() - await self.add_behavior(self._streaming_loop()) + self.add_behavior(self._streaming_loop()) # Start agents dependent on the output audio fragments here transcriber = TranscriptionAgent(audio_out_address) diff --git a/src/control_backend/core/agent_system.py b/src/control_backend/core/agent_system.py index ccdfe78..b1d8456 100644 --- a/src/control_backend/core/agent_system.py +++ b/src/control_backend/core/agent_system.py @@ -1,6 +1,7 @@ import asyncio import logging from abc import ABC, abstractmethod +from asyncio import Task from collections.abc import Coroutine import zmq @@ -75,8 +76,8 @@ class BaseAgent(ABC): await self.setup() # Start processing inbox and ZMQ messages - await self.add_behavior(self._process_inbox()) - await self.add_behavior(self._receive_internal_zmq_loop()) + self.add_behavior(self._process_inbox()) + self.add_behavior(self._receive_internal_zmq_loop()) async def stop(self): """Stops the agent.""" @@ -128,7 +129,7 @@ class BaseAgent(ABC): """Override this to handle incoming messages.""" raise NotImplementedError - async def add_behavior(self, coro: Coroutine): + def add_behavior(self, coro: Coroutine) -> Task: """ Helper to add a behavior to the agent. To add asynchronous behavior to an agent, define an `async` function and add it to the task list by calling :func:`add_behavior` @@ -138,3 +139,4 @@ class BaseAgent(ABC): task = asyncio.create_task(coro) self._tasks.add(task) task.add_done_callback(self._tasks.discard) + return task diff --git a/test/unit/agents/actuation/test_robot_speech_agent.py b/test/unit/agents/actuation/test_robot_speech_agent.py index 1ec2c6f..15324f6 100644 --- a/test/unit/agents/actuation/test_robot_speech_agent.py +++ b/test/unit/agents/actuation/test_robot_speech_agent.py @@ -25,24 +25,14 @@ async def test_setup_bind(zmq_context, mocker): settings = mocker.patch("control_backend.agents.actuation.robot_speech_agent.settings") settings.zmq_settings.internal_sub_address = "tcp://internal:1234" - # Swallow background task coroutines to avoid un-awaited warnings - class Swallow: - def __init__(self): - self.calls = 0 - - async def __call__(self, coro): - self.calls += 1 - coro.close() - - swallow = Swallow() - agent.add_behavior = swallow + agent.add_behavior = MagicMock() await agent.setup() fake_socket.bind.assert_any_call("tcp://localhost:5555") fake_socket.connect.assert_any_call("tcp://internal:1234") fake_socket.setsockopt.assert_any_call(zmq.SUBSCRIBE, b"command") - assert swallow.calls == 1 + agent.add_behavior.assert_called_once() @pytest.mark.asyncio @@ -53,22 +43,13 @@ async def test_setup_connect(zmq_context, mocker): settings = mocker.patch("control_backend.agents.actuation.robot_speech_agent.settings") settings.zmq_settings.internal_sub_address = "tcp://internal:1234" - class Swallow: - def __init__(self): - self.calls = 0 - - async def __call__(self, coro): - self.calls += 1 - coro.close() - - swallow = Swallow() - agent.add_behavior = swallow + agent.add_behavior = MagicMock() await agent.setup() fake_socket.connect.assert_any_call("tcp://localhost:5555") fake_socket.connect.assert_any_call("tcp://internal:1234") - assert swallow.calls == 1 + agent.add_behavior.assert_called_once() @pytest.mark.asyncio diff --git a/test/unit/agents/communication/test_ri_communication_agent.py b/test/unit/agents/communication/test_ri_communication_agent.py index 20b9379..747c4d2 100644 --- a/test/unit/agents/communication/test_ri_communication_agent.py +++ b/test/unit/agents/communication/test_ri_communication_agent.py @@ -46,16 +46,7 @@ async def test_setup_success_connects_and_starts_robot(zmq_context): robot_instance.start = AsyncMock() agent = RICommunicationAgent("ri_comm", address="tcp://localhost:5555", bind=False) - class Swallow: - def __init__(self): - self.calls = 0 - - async def __call__(self, coro): - self.calls += 1 - coro.close() - - swallow = Swallow() - agent.add_behavior = swallow + agent.add_behavior = MagicMock() await agent.setup() @@ -63,7 +54,7 @@ async def test_setup_success_connects_and_starts_robot(zmq_context): fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": {}}) robot_instance.start.assert_awaited_once() MockRobot.assert_called_once_with(ANY, address="tcp://*:5556", bind=True) - assert swallow.calls == 1 + agent.add_behavior.assert_called_once() assert agent.connected is True @@ -76,23 +67,14 @@ async def test_setup_binds_when_requested(zmq_context): agent = RICommunicationAgent("ri_comm", address="tcp://localhost:5555", bind=True) - class Swallow: - def __init__(self): - self.calls = 0 - - async def __call__(self, coro): - self.calls += 1 - coro.close() - - swallow = Swallow() - agent.add_behavior = swallow + agent.add_behavior = MagicMock() with patch(speech_agent_path(), autospec=True) as MockRobot: MockRobot.return_value.start = AsyncMock() await agent.setup() fake_socket.bind.assert_any_call("tcp://localhost:5555") - assert swallow.calls == 1 + agent.add_behavior.assert_called_once() @pytest.mark.asyncio diff --git a/test/unit/core/test_agent_system.py b/test/unit/core/test_agent_system.py index 5e954c8..f78b230 100644 --- a/test/unit/core/test_agent_system.py +++ b/test/unit/core/test_agent_system.py @@ -33,14 +33,16 @@ async def test_agent_lifecycle(): # Test background task async def dummy_task(): - await asyncio.sleep(0.01) + pass - await agent.add_behavior(dummy_task()) - assert len(agent._tasks) > 0 + task = agent.add_behavior(dummy_task()) + assert task in agent._tasks + + await task # Wait for task to finish - await asyncio.sleep(0.02) - assert len(agent._tasks) == 2 # message handling tasks are running + assert task not in agent._tasks + assert len(agent._tasks) == 2 # message handling tasks are still running await agent.stop() assert agent._running is False