109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
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
|