fix: merge refactor/zmq-internal-socket-behaviour into feat/cb2ui-robot-connections. (And fixed all ruff/ test issues to commit)
ref: None
This commit is contained in:
@@ -11,25 +11,27 @@ from control_backend.agents.ri_command_agent import RICommandAgent
|
||||
async def test_setup_bind(monkeypatch):
|
||||
"""Test setup with bind=True"""
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_command_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
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_comm_address="tcp://internal:1234")),
|
||||
MagicMock(zmq_settings=MagicMock(internal_sub_address="tcp://internal:1234")),
|
||||
)
|
||||
|
||||
await agent.setup()
|
||||
|
||||
# Ensure PUB socket bound
|
||||
fake_socket.bind.assert_any_call("tcp://localhost:5555")
|
||||
# Ensure SUB socket connected to internal address and subscribed
|
||||
fake_socket.connect.assert_any_call("tcp://internal:1234")
|
||||
fake_socket.setsockopt.assert_any_call(zmq.SUBSCRIBE, b"command")
|
||||
|
||||
# Ensure behaviour attached
|
||||
assert any(isinstance(b, agent.SendCommandsBehaviour) for b in agent.behaviours)
|
||||
|
||||
|
||||
@@ -37,19 +39,23 @@ async def test_setup_bind(monkeypatch):
|
||||
async def test_setup_connect(monkeypatch):
|
||||
"""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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_command_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
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_comm_address="tcp://internal:1234")),
|
||||
MagicMock(zmq_settings=MagicMock(internal_sub_address="tcp://internal:1234")),
|
||||
)
|
||||
|
||||
await agent.setup()
|
||||
|
||||
# Ensure PUB socket connected
|
||||
fake_socket.connect.assert_any_call("tcp://localhost:5555")
|
||||
|
||||
|
||||
|
||||
@@ -93,12 +93,14 @@ async def test_setup_creates_socket_and_negotiate_1(monkeypatch):
|
||||
fake_socket = MagicMock()
|
||||
fake_socket.send_json = AsyncMock()
|
||||
fake_socket.recv_json = fake_json_correct_negototiate_1()
|
||||
|
||||
fake_pub_socket = AsyncMock()
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
# Mock RICommandAgent agent startup
|
||||
@@ -107,13 +109,11 @@ 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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
@@ -143,10 +143,14 @@ async def test_setup_creates_socket_and_negotiate_2(monkeypatch):
|
||||
fake_socket = MagicMock()
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
# Mock RICommandAgent agent startup
|
||||
@@ -155,13 +159,11 @@ 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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
@@ -191,10 +193,14 @@ async def test_setup_creates_socket_and_negotiate_3(monkeypatch, caplog):
|
||||
fake_socket = MagicMock()
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
# Mock RICommandAgent agent startup
|
||||
@@ -206,13 +212,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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
@@ -239,10 +243,14 @@ async def test_setup_creates_socket_and_negotiate_4(monkeypatch):
|
||||
fake_socket = MagicMock()
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
# Mock RICommandAgent agent startup
|
||||
@@ -251,12 +259,10 @@ 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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=True,
|
||||
)
|
||||
@@ -286,10 +292,14 @@ async def test_setup_creates_socket_and_negotiate_5(monkeypatch):
|
||||
fake_socket = MagicMock()
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
# Mock RICommandAgent agent startup
|
||||
@@ -298,12 +308,10 @@ 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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
@@ -333,10 +341,14 @@ async def test_setup_creates_socket_and_negotiate_6(monkeypatch):
|
||||
fake_socket = MagicMock()
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
# Mock RICommandAgent agent startup
|
||||
@@ -345,12 +357,10 @@ 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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
@@ -380,10 +390,14 @@ async def test_setup_creates_socket_and_negotiate_7(monkeypatch, caplog):
|
||||
fake_socket = MagicMock()
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
# Mock RICommandAgent agent startup
|
||||
@@ -395,14 +409,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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
@@ -426,10 +438,14 @@ async def test_setup_creates_socket_and_negotiate_timeout(monkeypatch, caplog):
|
||||
fake_socket = MagicMock()
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
with patch(
|
||||
@@ -437,14 +453,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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
@@ -466,11 +480,11 @@ 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()
|
||||
fake_socket.send_multipart = AsyncMock()
|
||||
|
||||
# TODO: Integration test between actual server and password needed for spade agents
|
||||
agent = RICommunicationAgent("test@server", "password", fake_pub_socket)
|
||||
agent.req_socket = fake_socket
|
||||
agent = RICommunicationAgent("test@server", "password")
|
||||
agent._req_socket = fake_socket
|
||||
|
||||
behaviour = agent.ListenBehaviour()
|
||||
agent.add_behaviour(behaviour)
|
||||
@@ -505,7 +519,7 @@ async def test_listen_behaviour_ping_wrong_endpoint(caplog):
|
||||
fake_pub_socket = AsyncMock()
|
||||
|
||||
agent = RICommunicationAgent("test@server", "password", fake_pub_socket)
|
||||
agent.req_socket = fake_socket
|
||||
agent._req_socket = fake_socket
|
||||
|
||||
behaviour = agent.ListenBehaviour()
|
||||
agent.add_behaviour(behaviour)
|
||||
@@ -525,10 +539,10 @@ 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()
|
||||
fake_socket.send_multipart = AsyncMock()
|
||||
|
||||
agent = RICommunicationAgent("test@server", "password", fake_pub_socket)
|
||||
agent.req_socket = fake_socket
|
||||
agent = RICommunicationAgent("test@server", "password")
|
||||
agent._req_socket = fake_socket
|
||||
|
||||
behaviour = agent.ListenBehaviour()
|
||||
agent.add_behaviour(behaviour)
|
||||
@@ -546,6 +560,7 @@ async def test_listen_behaviour_ping_no_endpoint(caplog):
|
||||
"""
|
||||
fake_socket = AsyncMock()
|
||||
fake_socket.send_json = AsyncMock()
|
||||
fake_socket.send_multipart = AsyncMock()
|
||||
|
||||
# This is a message without endpoint >:(
|
||||
fake_socket.recv_json = AsyncMock(
|
||||
@@ -553,10 +568,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", fake_pub_socket)
|
||||
agent.req_socket = fake_socket
|
||||
agent = RICommunicationAgent("test@server", "password")
|
||||
agent._req_socket = fake_socket
|
||||
|
||||
behaviour = agent.ListenBehaviour()
|
||||
agent.add_behaviour(behaviour)
|
||||
@@ -574,18 +588,20 @@ 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!"))
|
||||
fake_socket.send_multipart = AsyncMock()
|
||||
|
||||
fake_context = MagicMock()
|
||||
fake_context.socket.return_value = fake_socket
|
||||
monkeypatch.setattr(
|
||||
"control_backend.agents.ri_communication_agent.context.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
agent = RICommunicationAgent(
|
||||
"test@server",
|
||||
"password",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
@@ -602,6 +618,7 @@ async def test_setup_unpacking_exception(monkeypatch, caplog):
|
||||
# --- Arrange ---
|
||||
fake_socket = MagicMock()
|
||||
fake_socket.send_json = AsyncMock()
|
||||
fake_socket.send_multipart = AsyncMock()
|
||||
|
||||
# Make recv_json return malformed negotiation data to trigger unpacking exception
|
||||
malformed_data = {
|
||||
@@ -611,8 +628,11 @@ async def test_setup_unpacking_exception(monkeypatch, caplog):
|
||||
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.socket", lambda _: fake_socket
|
||||
"control_backend.agents.ri_communication_agent.Context",
|
||||
MagicMock(instance=MagicMock(return_value=fake_context)),
|
||||
)
|
||||
|
||||
# Patch RICommandAgent so it won't actually start
|
||||
@@ -621,12 +641,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",
|
||||
pub_socket=fake_pub_socket,
|
||||
address="tcp://localhost:5555",
|
||||
bind=False,
|
||||
)
|
||||
|
||||
Binary file not shown.
108
test/integration/agents/vad_agent/test_vad_agent.py
Normal file
108
test/integration/agents/vad_agent/test_vad_agent.py
Normal file
@@ -0,0 +1,108 @@
|
||||
import random
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import zmq
|
||||
from spade.agent import Agent
|
||||
|
||||
from control_backend.agents.vad_agent import VADAgent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def zmq_context(mocker):
|
||||
return mocker.patch("control_backend.agents.vad_agent.zmq_context")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def streaming(mocker):
|
||||
return mocker.patch("control_backend.agents.vad_agent.Streaming")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def transcription_agent(mocker):
|
||||
return mocker.patch("control_backend.agents.vad_agent.TranscriptionAgent", autospec=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normal_setup(streaming, transcription_agent):
|
||||
"""
|
||||
Test that during normal setup, the VAD agent creates a Streaming behavior and creates audio
|
||||
sockets, and starts the TranscriptionAgent without loading real models.
|
||||
"""
|
||||
vad_agent = VADAgent("tcp://localhost:12345", False)
|
||||
vad_agent.add_behaviour = MagicMock()
|
||||
|
||||
await vad_agent.setup()
|
||||
|
||||
streaming.assert_called_once()
|
||||
vad_agent.add_behaviour.assert_called_once_with(streaming.return_value)
|
||||
transcription_agent.assert_called_once()
|
||||
transcription_agent.return_value.start.assert_called_once()
|
||||
assert vad_agent.audio_in_socket is not None
|
||||
assert vad_agent.audio_out_socket is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("do_bind", [True, False])
|
||||
def test_in_socket_creation(zmq_context, do_bind: bool):
|
||||
"""
|
||||
Test that the VAD agent creates an audio input socket, differentiating between binding and
|
||||
connecting.
|
||||
"""
|
||||
vad_agent = VADAgent(f"tcp://{'*' if do_bind else 'localhost'}:12345", do_bind)
|
||||
|
||||
vad_agent._connect_audio_in_socket()
|
||||
|
||||
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, "")
|
||||
|
||||
if do_bind:
|
||||
zmq_context.socket.return_value.bind.assert_called_once_with("tcp://*:12345")
|
||||
else:
|
||||
zmq_context.socket.return_value.connect.assert_called_once_with("tcp://localhost:12345")
|
||||
|
||||
|
||||
def test_out_socket_creation(zmq_context):
|
||||
"""
|
||||
Test that the VAD agent creates an audio output socket correctly.
|
||||
"""
|
||||
vad_agent = VADAgent("tcp://localhost:12345", False)
|
||||
|
||||
vad_agent._connect_audio_out_socket()
|
||||
|
||||
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()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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
|
||||
vad_agent = VADAgent("tcp://localhost:12345", False)
|
||||
|
||||
await vad_agent.setup()
|
||||
|
||||
assert vad_agent.audio_out_socket is None
|
||||
mock_super_stop.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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)
|
||||
|
||||
await vad_agent.setup()
|
||||
await vad_agent.stop()
|
||||
|
||||
assert zmq_context.socket.return_value.close.call_count == 2
|
||||
assert vad_agent.audio_in_socket is None
|
||||
assert vad_agent.audio_out_socket is None
|
||||
57
test/integration/agents/vad_agent/test_vad_with_audio.py
Normal file
57
test/integration/agents/vad_agent/test_vad_with_audio.py
Normal file
@@ -0,0 +1,57 @@
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
import soundfile as sf
|
||||
import zmq
|
||||
|
||||
from control_backend.agents.vad_agent import Streaming
|
||||
|
||||
|
||||
def get_audio_chunks() -> list[bytes]:
|
||||
curr_file = os.path.realpath(__file__)
|
||||
curr_dir = os.path.dirname(curr_file)
|
||||
file = f"{curr_dir}/speech_with_pauses_16k_1c_float32.wav"
|
||||
|
||||
chunk_size = 512
|
||||
|
||||
chunks = []
|
||||
|
||||
with sf.SoundFile(file, "r") as f:
|
||||
assert f.samplerate == 16000
|
||||
assert f.channels == 1
|
||||
assert f.subtype == "FLOAT"
|
||||
|
||||
while True:
|
||||
data = f.read(chunk_size, dtype="float32")
|
||||
if len(data) != chunk_size:
|
||||
break
|
||||
|
||||
chunks.append(data.tobytes())
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_audio(mocker):
|
||||
"""
|
||||
Test the VAD agent with only input and output mocked. Using the real model, using real audio as
|
||||
input. Ensure that it outputs some fragments with audio.
|
||||
"""
|
||||
audio_chunks = get_audio_chunks()
|
||||
audio_in_socket = AsyncMock()
|
||||
audio_in_socket.recv.side_effect = audio_chunks
|
||||
|
||||
mock_poller: MagicMock = mocker.patch("control_backend.agents.vad_agent.zmq.Poller")
|
||||
mock_poller.return_value.poll.return_value = [(audio_in_socket, zmq.POLLIN)]
|
||||
|
||||
audio_out_socket = AsyncMock()
|
||||
|
||||
vad_streamer = Streaming(audio_in_socket, audio_out_socket)
|
||||
for _ in audio_chunks:
|
||||
await vad_streamer.run()
|
||||
|
||||
audio_out_socket.send.assert_called()
|
||||
for args in audio_out_socket.send.call_args_list:
|
||||
assert isinstance(args[0][0], bytes)
|
||||
assert len(args[0][0]) >= 512 * 4 * 3 # Should be at least 3 chunks of audio
|
||||
Reference in New Issue
Block a user