fix: fixed new tests and merged dev into branch

ref: N25B-256

test/integration/agents/test_ri_commands_agent.py

@@ -7,25 +7,21 @@ import zmq
 from control_backend.agents.ri_command_agent import RICommandAgent
 
 
-@pytest.mark.asyncio
-async def test_setup_bind(monkeypatch):
-    """Test setup with bind=True"""
-    fake_socket = MagicMock()
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
+@pytest.fixture
+def zmq_context(mocker):
+    mock_context = mocker.patch("control_backend.agents.vad_agent.azmq.Context.instance")
+    mock_context.return_value = MagicMock()
+    return mock_context
 
-    # Patch Context.instance() to return fake_context
-    monkeypatch.setattr(
-        "control_backend.agents.ri_command_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
+
+@pytest.mark.asyncio
+async def test_setup_bind(zmq_context, mocker):
+    """Test setup with bind=True"""
+    fake_socket = zmq_context.return_value.socket.return_value
 
     agent = RICommandAgent("test@server", "password", address="tcp://localhost:5555", bind=True)
-
-    monkeypatch.setattr(
-        "control_backend.agents.ri_command_agent.settings",
-        MagicMock(zmq_settings=MagicMock(internal_sub_address="tcp://internal:1234")),
-    )
+    settings = mocker.patch("control_backend.agents.ri_command_agent.settings")
+    settings.zmq_settings.internal_sub_address = "tcp://internal:1234"
 
     await agent.setup()
 
@@ -36,23 +32,13 @@ async def test_setup_bind(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_setup_connect(monkeypatch):
+async def test_setup_connect(zmq_context, mocker):
     """Test setup with bind=False"""
-    fake_socket = MagicMock()
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-
-    # Patch Context.instance() to return fake_context
-    monkeypatch.setattr(
-        "control_backend.agents.ri_command_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
+    fake_socket = zmq_context.return_value.socket.return_value
 
     agent = RICommandAgent("test@server", "password", address="tcp://localhost:5555", bind=False)
-    monkeypatch.setattr(
-        "control_backend.agents.ri_command_agent.settings",
-        MagicMock(zmq_settings=MagicMock(internal_sub_address="tcp://internal:1234")),
-    )
+    settings = mocker.patch("control_backend.agents.ri_command_agent.settings")
+    settings.zmq_settings.internal_sub_address = "tcp://internal:1234"
 
     await agent.setup()
 

test/integration/agents/test_ri_communication_agent.py

@@ -84,25 +84,24 @@ def fake_json_invalid_id_negototiate():
     )
 
 
+@pytest.fixture
+def zmq_context(mocker):
+    mock_context = mocker.patch("control_backend.agents.vad_agent.azmq.Context.instance")
+    mock_context.return_value = MagicMock()
+    return mock_context
+
+
 @pytest.mark.asyncio
-async def test_setup_creates_socket_and_negotiate_1(monkeypatch):
+async def test_setup_creates_socket_and_negotiate_1(zmq_context):
     """
     Test the setup of the communication agent
     """
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = fake_json_correct_negototiate_1()
     fake_socket.send_multipart = AsyncMock()
 
-    # Mock context.socket to return our fake socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     # Mock RICommandAgent agent startup
     with patch(
         "control_backend.agents.ri_communication_agent.RICommandAgent", autospec=True
@@ -135,24 +134,16 @@ async def test_setup_creates_socket_and_negotiate_1(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_setup_creates_socket_and_negotiate_2(monkeypatch):
+async def test_setup_creates_socket_and_negotiate_2(zmq_context):
     """
     Test the setup of the communication agent
     """
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = fake_json_correct_negototiate_2()
     fake_socket.send_multipart = AsyncMock()
 
-    # Mock context.socket to return our fake socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     # Mock RICommandAgent agent startup
     with patch(
         "control_backend.agents.ri_communication_agent.RICommandAgent", autospec=True
@@ -185,24 +176,16 @@ async def test_setup_creates_socket_and_negotiate_2(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_setup_creates_socket_and_negotiate_3(monkeypatch, caplog):
+async def test_setup_creates_socket_and_negotiate_3(zmq_context, caplog):
     """
     Test the functionality of setup with incorrect negotiation message
     """
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = fake_json_wrong_negototiate_1()
     fake_socket.send_multipart = AsyncMock()
 
-    # Mock context.socket to return our fake socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     # Mock RICommandAgent agent startup
 
     # We are sending wrong negotiation info to the communication agent,
@@ -235,24 +218,16 @@ async def test_setup_creates_socket_and_negotiate_3(monkeypatch, caplog):
 
 
 @pytest.mark.asyncio
-async def test_setup_creates_socket_and_negotiate_4(monkeypatch):
+async def test_setup_creates_socket_and_negotiate_4(zmq_context):
     """
     Test the setup of the communication agent with different bind value
     """
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = fake_json_correct_negototiate_3()
     fake_socket.send_multipart = AsyncMock()
 
-    # Mock context.socket to return our fake socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     # Mock RICommandAgent agent startup
     with patch(
         "control_backend.agents.ri_communication_agent.RICommandAgent", autospec=True
@@ -284,24 +259,16 @@ async def test_setup_creates_socket_and_negotiate_4(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_setup_creates_socket_and_negotiate_5(monkeypatch):
+async def test_setup_creates_socket_and_negotiate_5(zmq_context):
     """
     Test the setup of the communication agent
     """
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = fake_json_correct_negototiate_4()
     fake_socket.send_multipart = AsyncMock()
 
-    # Mock context.socket to return our fake socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     # Mock RICommandAgent agent startup
     with patch(
         "control_backend.agents.ri_communication_agent.RICommandAgent", autospec=True
@@ -333,24 +300,16 @@ async def test_setup_creates_socket_and_negotiate_5(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_setup_creates_socket_and_negotiate_6(monkeypatch):
+async def test_setup_creates_socket_and_negotiate_6(zmq_context):
     """
     Test the setup of the communication agent
     """
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = fake_json_correct_negototiate_5()
     fake_socket.send_multipart = AsyncMock()
 
-    # Mock context.socket to return our fake socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     # Mock RICommandAgent agent startup
     with patch(
         "control_backend.agents.ri_communication_agent.RICommandAgent", autospec=True
@@ -382,28 +341,20 @@ async def test_setup_creates_socket_and_negotiate_6(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_setup_creates_socket_and_negotiate_7(monkeypatch, caplog):
+async def test_setup_creates_socket_and_negotiate_7(zmq_context, caplog):
     """
     Test the functionality of setup with incorrect id
     """
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = fake_json_invalid_id_negototiate()
     fake_socket.send_multipart = AsyncMock()
 
-    # Mock context.socket to return our fake socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     # Mock RICommandAgent agent startup
 
     # We are sending wrong negotiation info to the communication agent,
-    # so we should retry and expect a etter response, within a limited time.
+    # so we should retry and expect a better response, within a limited time.
     with patch(
         "control_backend.agents.ri_communication_agent.RICommandAgent", autospec=True
     ) as MockCommandAgent:
@@ -430,24 +381,16 @@ async def test_setup_creates_socket_and_negotiate_7(monkeypatch, caplog):
 
 
 @pytest.mark.asyncio
-async def test_setup_creates_socket_and_negotiate_timeout(monkeypatch, caplog):
+async def test_setup_creates_socket_and_negotiate_timeout(zmq_context, caplog):
     """
     Test the functionality of setup with incorrect negotiation message
     """
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.recv_json = AsyncMock(side_effect=asyncio.TimeoutError)
     fake_socket.send_multipart = AsyncMock()
 
-    # Mock context.socket to return our fake socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     with patch(
         "control_backend.agents.ri_communication_agent.RICommandAgent", autospec=True
     ) as MockCommandAgent:
@@ -534,8 +477,8 @@ async def test_listen_behaviour_ping_wrong_endpoint(caplog):
 
 
 @pytest.mark.asyncio
-async def test_listen_behaviour_timeout(caplog):
-    fake_socket = AsyncMock()
+async def test_listen_behaviour_timeout(zmq_context, caplog):
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     # recv_json will never resolve, simulate timeout
     fake_socket.recv_json = AsyncMock(side_effect=asyncio.TimeoutError)
@@ -585,20 +528,13 @@ async def test_listen_behaviour_ping_no_endpoint(caplog):
 
 
 @pytest.mark.asyncio
-async def test_setup_unexpected_exception(monkeypatch, caplog):
-    fake_socket = MagicMock()
+async def test_setup_unexpected_exception(zmq_context, caplog):
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     # Simulate unexpected exception during recv_json()
     fake_socket.recv_json = AsyncMock(side_effect=Exception("boom!"))
     fake_socket.send_multipart = AsyncMock()
 
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     agent = RICommunicationAgent(
         "test@server",
         "password",
@@ -614,9 +550,9 @@ async def test_setup_unexpected_exception(monkeypatch, caplog):
 
 
 @pytest.mark.asyncio
-async def test_setup_unpacking_exception(monkeypatch, caplog):
+async def test_setup_unpacking_exception(zmq_context, caplog):
     # --- Arrange ---
-    fake_socket = MagicMock()
+    fake_socket = zmq_context.return_value.socket.return_value
     fake_socket.send_json = AsyncMock()
     fake_socket.send_multipart = AsyncMock()
 
@@ -627,14 +563,6 @@ async def test_setup_unpacking_exception(monkeypatch, caplog):
     }  # missing 'port' and 'bind'
     fake_socket.recv_json = AsyncMock(return_value=malformed_data)
 
-    # Patch context.socket
-    fake_context = MagicMock()
-    fake_context.socket.return_value = fake_socket
-    monkeypatch.setattr(
-        "control_backend.agents.ri_communication_agent.Context",
-        MagicMock(instance=MagicMock(return_value=fake_context)),
-    )
-
     # Patch RICommandAgent so it won't actually start
     with patch(
         "control_backend.agents.ri_communication_agent.RICommandAgent", autospec=True

test/integration/agents/vad_agent/test_vad_agent.py

@@ -10,7 +10,9 @@ from control_backend.agents.vad_agent import VADAgent
 
 @pytest.fixture
 def zmq_context(mocker):
-    return mocker.patch("control_backend.agents.vad_agent.zmq_context")
+    mock_context = mocker.patch("control_backend.agents.vad_agent.azmq.Context.instance")
+    mock_context.return_value = MagicMock()
+    return mock_context
 
 
 @pytest.fixture
@@ -54,13 +56,18 @@ def test_in_socket_creation(zmq_context, do_bind: bool):
 
     assert vad_agent.audio_in_socket is not None
 
-    zmq_context.socket.assert_called_once_with(zmq.SUB)
-    zmq_context.socket.return_value.setsockopt_string.assert_called_once_with(zmq.SUBSCRIBE, "")
+    zmq_context.return_value.socket.assert_called_once_with(zmq.SUB)
+    zmq_context.return_value.socket.return_value.setsockopt_string.assert_called_once_with(
+        zmq.SUBSCRIBE,
+        "",
+    )
 
     if do_bind:
-        zmq_context.socket.return_value.bind.assert_called_once_with("tcp://*:12345")
+        zmq_context.return_value.socket.return_value.bind.assert_called_once_with("tcp://*:12345")
     else:
-        zmq_context.socket.return_value.connect.assert_called_once_with("tcp://localhost:12345")
+        zmq_context.return_value.socket.return_value.connect.assert_called_once_with(
+            "tcp://localhost:12345"
+        )
 
 
 def test_out_socket_creation(zmq_context):
@@ -73,8 +80,8 @@ def test_out_socket_creation(zmq_context):
 
     assert vad_agent.audio_out_socket is not None
 
-    zmq_context.socket.assert_called_once_with(zmq.PUB)
-    zmq_context.socket.return_value.bind_to_random_port.assert_called_once()
+    zmq_context.return_value.socket.assert_called_once_with(zmq.PUB)
+    zmq_context.return_value.socket.return_value.bind_to_random_port.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -83,7 +90,9 @@ async def test_out_socket_creation_failure(zmq_context):
     Test setup failure when the audio output socket cannot be created.
     """
     with patch.object(Agent, "stop", new_callable=AsyncMock) as mock_super_stop:
-        zmq_context.socket.return_value.bind_to_random_port.side_effect = zmq.ZMQBindError
+        zmq_context.return_value.socket.return_value.bind_to_random_port.side_effect = (
+            zmq.ZMQBindError
+        )
         vad_agent = VADAgent("tcp://localhost:12345", False)
 
         await vad_agent.setup()
@@ -98,11 +107,14 @@ async def test_stop(zmq_context, transcription_agent):
     Test that when the VAD agent is stopped, the sockets are closed correctly.
     """
     vad_agent = VADAgent("tcp://localhost:12345", False)
-    zmq_context.socket.return_value.bind_to_random_port.return_value = random.randint(1000, 10000)
+    zmq_context.return_value.socket.return_value.bind_to_random_port.return_value = random.randint(
+        1000,
+        10000,
+    )
 
     await vad_agent.setup()
     await vad_agent.stop()
 
-    assert zmq_context.socket.return_value.close.call_count == 2
+    assert zmq_context.return_value.socket.return_value.close.call_count == 2
     assert vad_agent.audio_in_socket is None
     assert vad_agent.audio_out_socket is None

test/integration/agents/vad_agent/test_vad_with_audio.py

@@ -48,6 +48,7 @@ async def test_real_audio(mocker):
     audio_out_socket = AsyncMock()
 
     vad_streamer = Streaming(audio_in_socket, audio_out_socket)
+    vad_streamer._ready = True
     for _ in audio_chunks:
         await vad_streamer.run()
 

test/integration/api/endpoints/test_robot_endpoint.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock
 
 import pytest
 from fastapi import FastAPI
@@ -16,7 +16,6 @@ def app():
     """
     app = FastAPI()
     app.include_router(robot.router)
-    app.state.internal_comm_socket = MagicMock()  # mock ZMQ socket
     return app
 
 
@@ -26,32 +25,30 @@ def client(app):
     return TestClient(app)
 
 
-def test_receive_command_endpoint(client, app):
+def test_receive_command_success(client):
     """
-    Test that a POST to /command sends the right multipart message
-    and returns a 202 with the expected JSON body.
+    Test for successful reception of a command. Ensures the status code is 202 and the response body
+    is correct. It also verifies that the ZeroMQ socket's send_multipart method is called with the
+    expected data.
     """
-    mock_socket = app.state.internal_comm_socket
+    # Arrange
+    mock_pub_socket = AsyncMock()
+    client.app.state.endpoints_pub_socket = mock_pub_socket
 
-    # Prepare test payload that matches SpeechCommand
-    payload = {"endpoint": "actuate/speech", "data": "yooo"}
+    command_data = {"endpoint": "actuate/speech", "data": "This is a test"}
+    speech_command = SpeechCommand(**command_data)
 
-    # Send POST request
-    response = client.post("/command", json=payload)
+    # Act
+    response = client.post("/command", json=command_data)
 
-    # Check response
+    # Assert
     assert response.status_code == 202
     assert response.json() == {"status": "Command received"}
 
-    # Verify that the socket was called with the correct data
-    assert mock_socket.send_multipart.called, "Socket should be used to send data"
-
-    args, kwargs = mock_socket.send_multipart.call_args
-    sent_data = args[0]
-
-    assert sent_data[0] == b"command"
-    # Check JSON encoding roughly matches
-    assert isinstance(SpeechCommand.model_validate_json(sent_data[1].decode()), SpeechCommand)
+    # Verify that the ZMQ socket was used correctly
+    mock_pub_socket.send_multipart.assert_awaited_once_with(
+        [b"command", speech_command.model_dump_json().encode()]
+    )
 
 
 def test_receive_command_invalid_payload(client):

test/integration/schemas/test_ri_message.py

@@ -16,12 +16,11 @@ def test_valid_speech_command_1():
     command = valid_command_1()
     RIMessage.model_validate(command)
     SpeechCommand.model_validate(command)
-    assert True
 
 
 def test_invalid_speech_command_1():
     command = invalid_command_1()
     RIMessage.model_validate(command)
+
     with pytest.raises(ValidationError):
         SpeechCommand.model_validate(command)
-    assert True

test/unit/agents/belief_collector/behaviours/test_continuous_collect.py

@@ -182,8 +182,6 @@ async def test_belief_text_values_not_lists(continuous_collector, mocker):
 @pytest.mark.asyncio
 async def test_belief_text_happy_path_logs_items_and_sends(continuous_collector, mocker):
     payload = {"type": "belief_extraction_text", "beliefs": {"user_said": ["hello test", "No"]}}
-    # Your code calls self.send(..); patch it
-    # (or switch implementation to self.agent.send and patch that)
     continuous_collector.send = AsyncMock()
     logger_mock = mocker.patch(
         "control_backend.agents.belief_collector.behaviours.continuous_collect.logger"

test/unit/agents/test_vad_streaming.py

@@ -21,7 +21,9 @@ def streaming(audio_in_socket, audio_out_socket):
     import torch
 
     torch.hub.load.return_value = (..., ...)  # Mock
-    return Streaming(audio_in_socket, audio_out_socket)
+    streaming = Streaming(audio_in_socket, audio_out_socket)
+    streaming._ready = True
+    return streaming
 
 
 async def simulate_streaming_with_probabilities(streaming, probabilities: list[float]):

test/unit/agents/transcription/test_speech_recognizer.py

@@ -5,12 +5,13 @@ from control_backend.agents.transcription.speech_recognizer import OpenAIWhisper
 
 
 def test_estimate_max_tokens():
-    """Inputting one minute of audio, assuming 300 words per minute, expecting 400 tokens."""
+    """Inputting one minute of audio, assuming 450 words per minute and adding a 10 token padding,
+    expecting 610 tokens."""
     audio = np.empty(shape=(60 * 16_000), dtype=np.float32)
 
     actual = SpeechRecognizer._estimate_max_tokens(audio)
 
-    assert actual == 400
+    assert actual == 610
     assert isinstance(actual, int)