fix: fix test race condition

ref: N25B-301

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