fix: fix test race condition

ref: N25B-301

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)
 

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.")
 

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.")
 

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)
 

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)

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

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

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

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