59 lines
1.6 KiB
Python
59 lines
1.6 KiB
Python
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)
|
|
vad_streamer._ready = True
|
|
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
|