from unittest.mock import Mock, sentinel

import pytest
import zmq

import gitubuntu.importer_service_worker as target


def test_zmq_recv_with_timeout():
    """If data is available it should get returned"""
    poller = Mock(
        spec=zmq.Poller,
        poll=Mock(return_value=True),
    )
    socket = Mock(
        spec=zmq.Socket,
        recv=Mock(return_value=sentinel.result),
    )
    result = target.zmq_recv_with_timeout(poller, socket, None)
    assert result is sentinel.result


def test_zmq_recv_with_timeout_timeout():
    """If data is not available after timeout an exception should get raised"""
    poller = Mock(
        spec=zmq.Poller,
        poll=Mock(return_value=False),
    )
    with pytest.raises(target.BrokerTimeout):
        target.zmq_recv_with_timeout(poller, None, None)


def test_zmq_recv_with_timeout_multiple():
    """If multiple polls are required the whole frame should be returned"""
    poller = Mock(
        spec=zmq.Poller,
        poll=Mock(return_value=True),
    )
    socket = Mock(
        spec=zmq.Socket,
        recv=Mock(side_effect=[zmq.error.Again, sentinel.result]),
    )
    result = target.zmq_recv_with_timeout(poller, socket, None)
    assert result is sentinel.result
