fix: fix test race condition

ref: N25B-301

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