fix: fix unit tests with new feat

ref: N25B-151

test/integration/agents/test_ri_communication_agent.py

@@ -92,6 +92,8 @@ async def test_setup_creates_socket_and_negotiate_1(monkeypatch):
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = fake_json_correct_negototiate_1()
 
+    fake_pub_socket = AsyncMock()
+
     # Mock context.socket to return our fake socket
     monkeypatch.setattr(
         "control_backend.agents.ri_communication_agent.context.socket", lambda _: fake_socket
@@ -103,16 +105,17 @@ async def test_setup_creates_socket_and_negotiate_1(monkeypatch):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
+        fake_pub_socket = AsyncMock()
 
         # --- Act ---
         agent = RICommunicationAgent(
-            "test@server", "password", address="tcp://localhost:5555", bind=False
+            "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
         )
         await agent.setup()
 
         # --- Assert ---
         fake_socket.connect.assert_any_call("tcp://localhost:5555")
-        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": None})
+        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": {}})
         fake_socket.recv_json.assert_awaited()
         fake_agent_instance.start.assert_awaited()
         MockCommandAgent.assert_called_once_with(
@@ -146,16 +149,17 @@ async def test_setup_creates_socket_and_negotiate_2(monkeypatch):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
+        fake_pub_socket = AsyncMock()
 
         # --- Act ---
         agent = RICommunicationAgent(
-            "test@server", "password", address="tcp://localhost:5555", bind=False
+            "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
         )
         await agent.setup()
 
         # --- Assert ---
         fake_socket.connect.assert_any_call("tcp://localhost:5555")
-        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": None})
+        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": {}})
         fake_socket.recv_json.assert_awaited()
         fake_agent_instance.start.assert_awaited()
         MockCommandAgent.assert_called_once_with(
@@ -192,11 +196,11 @@ async def test_setup_creates_socket_and_negotiate_3(monkeypatch, caplog):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
-
+        fake_pub_socket = AsyncMock()
         # --- Act ---
         with caplog.at_level("ERROR"):
             agent = RICommunicationAgent(
-                "test@server", "password", address="tcp://localhost:5555", bind=False
+                "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
             )
             await agent.setup(max_retries=1)
 
@@ -233,16 +237,16 @@ async def test_setup_creates_socket_and_negotiate_4(monkeypatch):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
-
+        fake_pub_socket = AsyncMock()
         # --- Act ---
         agent = RICommunicationAgent(
-            "test@server", "password", address="tcp://localhost:5555", bind=True
+            "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=True
         )
         await agent.setup()
 
         # --- Assert ---
         fake_socket.bind.assert_any_call("tcp://localhost:5555")
-        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": None})
+        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": {}})
         fake_socket.recv_json.assert_awaited()
         fake_agent_instance.start.assert_awaited()
         MockCommandAgent.assert_called_once_with(
@@ -276,16 +280,16 @@ async def test_setup_creates_socket_and_negotiate_5(monkeypatch):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
-
+        fake_pub_socket = AsyncMock()
         # --- Act ---
         agent = RICommunicationAgent(
-            "test@server", "password", address="tcp://localhost:5555", bind=False
+            "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
         )
         await agent.setup()
 
         # --- Assert ---
         fake_socket.connect.assert_any_call("tcp://localhost:5555")
-        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": None})
+        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": {}})
         fake_socket.recv_json.assert_awaited()
         fake_agent_instance.start.assert_awaited()
         MockCommandAgent.assert_called_once_with(
@@ -319,16 +323,16 @@ async def test_setup_creates_socket_and_negotiate_6(monkeypatch):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
-
+        fake_pub_socket = AsyncMock()
         # --- Act ---
         agent = RICommunicationAgent(
-            "test@server", "password", address="tcp://localhost:5555", bind=False
+            "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
         )
         await agent.setup()
 
         # --- Assert ---
         fake_socket.connect.assert_any_call("tcp://localhost:5555")
-        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": None})
+        fake_socket.send_json.assert_any_call({"endpoint": "negotiate/ports", "data": {}})
         fake_socket.recv_json.assert_awaited()
         fake_agent_instance.start.assert_awaited()
         MockCommandAgent.assert_called_once_with(
@@ -365,11 +369,12 @@ async def test_setup_creates_socket_and_negotiate_7(monkeypatch, caplog):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
+        fake_pub_socket = AsyncMock()
 
         # --- Act ---
         with caplog.at_level("WARNING"):
             agent = RICommunicationAgent(
-                "test@server", "password", address="tcp://localhost:5555", bind=False
+                "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
             )
             await agent.setup(max_retries=1)
 
@@ -402,11 +407,12 @@ async def test_setup_creates_socket_and_negotiate_timeout(monkeypatch, caplog):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
+        fake_pub_socket = AsyncMock()
     
         # --- Act ---
         with caplog.at_level("WARNING"):
             agent = RICommunicationAgent(
-                "test@server", "password", address="tcp://localhost:5555", bind=False
+                "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
             )
             await agent.setup(max_retries=1)
 
@@ -426,9 +432,10 @@ async def test_listen_behaviour_ping_correct(caplog):
     fake_socket = AsyncMock()
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = AsyncMock(return_value={"endpoint": "ping", "data": {}})
+    fake_pub_socket = AsyncMock()
 
     # TODO: Integration test between actual server and password needed for spade agents
-    agent = RICommunicationAgent("test@server", "password")
+    agent = RICommunicationAgent("test@server", "password", fake_pub_socket)
     agent.req_socket = fake_socket
 
     behaviour = agent.ListenBehaviour()
@@ -461,8 +468,9 @@ async def test_listen_behaviour_ping_wrong_endpoint(caplog):
             ],
         }
     )
+    fake_pub_socket = AsyncMock()
 
-    agent = RICommunicationAgent("test@server", "password")
+    agent = RICommunicationAgent("test@server", "password", fake_pub_socket)
     agent.req_socket = fake_socket
 
     behaviour = agent.ListenBehaviour()
@@ -483,8 +491,9 @@ async def test_listen_behaviour_timeout(caplog):
     fake_socket.send_json = AsyncMock()
     # recv_json will never resolve, simulate timeout
     fake_socket.recv_json = AsyncMock(side_effect=asyncio.TimeoutError)
+    fake_pub_socket = AsyncMock()
 
-    agent = RICommunicationAgent("test@server", "password")
+    agent = RICommunicationAgent("test@server", "password", fake_pub_socket)
     agent.req_socket = fake_socket
 
     behaviour = agent.ListenBehaviour()
@@ -510,8 +519,9 @@ async def test_listen_behaviour_ping_no_endpoint(caplog):
             "data": "I dont have an endpoint >:)",
         }
     )
+    fake_pub_socket = AsyncMock()
 
-    agent = RICommunicationAgent("test@server", "password")
+    agent = RICommunicationAgent("test@server", "password", fake_pub_socket)
     agent.req_socket = fake_socket
 
     behaviour = agent.ListenBehaviour()
@@ -530,15 +540,17 @@ async def test_listen_behaviour_ping_no_endpoint(caplog):
 async def test_setup_unexpected_exception(monkeypatch, caplog):
     fake_socket = MagicMock()
     fake_socket.send_json = AsyncMock()
+    fake_pub_socket = AsyncMock()
     # Simulate unexpected exception during recv_json()
     fake_socket.recv_json = AsyncMock(side_effect=Exception("boom!"))
 
+
     monkeypatch.setattr(
         "control_backend.agents.ri_communication_agent.context.socket", lambda _: fake_socket
     )
 
     agent = RICommunicationAgent(
-        "test@server", "password", address="tcp://localhost:5555", bind=False
+        "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
     )
 
     with caplog.at_level("ERROR"):
@@ -572,9 +584,10 @@ async def test_setup_unpacking_exception(monkeypatch, caplog):
     ) as MockCommandAgent:
         fake_agent_instance = MockCommandAgent.return_value
         fake_agent_instance.start = AsyncMock()
+        fake_pub_socket = AsyncMock()
 
         agent = RICommunicationAgent(
-            "test@server", "password", address="tcp://localhost:5555", bind=False
+            "test@server", "password", pub_socket=fake_pub_socket, address="tcp://localhost:5555", bind=False
         )
 
         # --- Act & Assert ---