47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
import zmq
|
|
|
|
from control_backend.agents.vad_agent import SocketPoller
|
|
|
|
|
|
@pytest.fixture
|
|
def socket():
|
|
return AsyncMock()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_socket_poller_with_data(socket, mocker):
|
|
socket_data = b"test"
|
|
socket.recv.return_value = socket_data
|
|
|
|
mock_poller: MagicMock = mocker.patch("control_backend.agents.vad_agent.zmq.Poller")
|
|
mock_poller.return_value.poll.return_value = [(socket, zmq.POLLIN)]
|
|
|
|
poller = SocketPoller(socket)
|
|
# Calling `poll` twice to be able to check that the poller is reused
|
|
await poller.poll()
|
|
data = await poller.poll()
|
|
|
|
assert data == socket_data
|
|
|
|
# Ensure that the poller was reused
|
|
mock_poller.assert_called_once_with()
|
|
mock_poller.return_value.register.assert_called_once_with(socket, zmq.POLLIN)
|
|
|
|
assert socket.recv.call_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_socket_poller_no_data(socket, mocker):
|
|
mock_poller: MagicMock = mocker.patch("control_backend.agents.vad_agent.zmq.Poller")
|
|
mock_poller.return_value.poll.return_value = []
|
|
|
|
poller = SocketPoller(socket)
|
|
data = await poller.poll()
|
|
|
|
assert data is None
|
|
|
|
socket.recv.assert_not_called()
|