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()