path: root/t/unit/test_pidbox.py

import socket
import warnings
from unittest.mock import Mock, patch
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor

import pytest

from kombu import Connection, pidbox
from kombu.exceptions import ContentDisallowed, InconsistencyError
from kombu.utils.uuid import uuid


def is_cast(message):
    return message['method']


def is_call(message):
    return message['method'] and message['reply_to']


class test_Mailbox:

    class Mailbox(pidbox.Mailbox):

        def _collect(self, *args, **kwargs):
            return 'COLLECTED'

    def setup(self):
        self.mailbox = self.Mailbox('test_pidbox')
        self.connection = Connection(transport='memory')
        self.state = {'var': 1}
        self.handlers = {'mymethod': self._handler}
        self.bound = self.mailbox(self.connection)
        self.default_chan = self.connection.channel()
        self.node = self.bound.Node(
            'test_pidbox',
            state=self.state, handlers=self.handlers,
            channel=self.default_chan,
        )

    def _handler(self, state):
        return self.stats['var']

    def test_broadcast_matcher_pattern_string_type(self):
        mailbox = pidbox.Mailbox("test_matcher_str")(self.connection)
        with pytest.raises(ValueError):
            mailbox._broadcast("ping", pattern=1, matcher=2)

    def test_publish_reply_ignores_InconsistencyError(self):
        mailbox = pidbox.Mailbox('test_reply__collect')(self.connection)
        with patch('kombu.pidbox.Producer') as Producer:
            producer = Producer.return_value = Mock(name='producer')
            producer.publish.side_effect = InconsistencyError()
            mailbox._publish_reply(
                {'foo': 'bar'}, mailbox.reply_exchange, mailbox.oid, 'foo',
            )
            producer.publish.assert_called()

    def test_reply__collect(self):
        mailbox = pidbox.Mailbox('test_reply__collect')(self.connection)
        exchange = mailbox.reply_exchange.name
        channel = self.connection.channel()
        mailbox.reply_queue(channel).declare()

        ticket = uuid()
        mailbox._publish_reply({'foo': 'bar'}, exchange, mailbox.oid, ticket)
        _callback_called = [False]

        def callback(body):
            _callback_called[0] = True

        reply = mailbox._collect(ticket, limit=1,
                                 callback=callback, channel=channel)
        assert reply == [{'foo': 'bar'}]
        assert _callback_called[0]

        ticket = uuid()
        mailbox._publish_reply({'biz': 'boz'}, exchange, mailbox.oid, ticket)
        reply = mailbox._collect(ticket, limit=1, channel=channel)
        assert reply == [{'biz': 'boz'}]

        mailbox._publish_reply({'foo': 'BAM'}, exchange, mailbox.oid, 'doom',
                               serializer='pickle')
        with pytest.raises(ContentDisallowed):
            reply = mailbox._collect('doom', limit=1, channel=channel)
        mailbox._publish_reply(
            {'foo': 'BAMBAM'}, exchange, mailbox.oid, 'doom',
            serializer='pickle',
        )
        reply = mailbox._collect('doom', limit=1, channel=channel,
                                 accept=['pickle'])
        assert reply[0]['foo'] == 'BAMBAM'

        de = mailbox.connection.drain_events = Mock()
        de.side_effect = socket.timeout
        mailbox._collect(ticket, limit=1, channel=channel)

    def test_reply__collect_uses_default_channel(self):
        class ConsumerCalled(Exception):
            pass

        def fake_Consumer(channel, *args, **kwargs):
            raise ConsumerCalled(channel)

        ticket = uuid()
        with patch('kombu.pidbox.Consumer') as Consumer:
            mailbox = pidbox.Mailbox('test_reply__collect')(self.connection)
            assert mailbox.connection.default_channel is not None
            Consumer.side_effect = fake_Consumer
            try:
                mailbox._collect(ticket, limit=1)
            except ConsumerCalled as c:
                assert c.args[0] is not None
            except Exception:
                raise
            else:
                assert False, "Consumer not called"

    def test__publish_uses_default_channel(self):
        class QueueCalled(Exception):
            pass

        def queue__call__side(channel, *args, **kwargs):
            raise QueueCalled(channel)

        ticket = uuid()
        with patch.object(pidbox.Queue, '__call__') as queue__call__:
            mailbox = pidbox.Mailbox('test_reply__collect')(self.connection)
            queue__call__.side_effect = queue__call__side
            try:
                mailbox._publish(ticket, {}, reply_ticket=ticket)
            except QueueCalled as c:
                assert c.args[0] is not None
            except Exception:
                raise
            else:
                assert False, "Queue not called"

    def test_constructor(self):
        assert self.mailbox.connection is None
        assert self.mailbox.exchange.name
        assert self.mailbox.reply_exchange.name

    def test_bound(self):
        bound = self.mailbox(self.connection)
        assert bound.connection is self.connection

    def test_Node(self):
        assert self.node.hostname
        assert self.node.state
        assert self.node.mailbox is self.bound
        assert self.handlers

        # No initial handlers
        node2 = self.bound.Node('test_pidbox2', state=self.state)
        assert node2.handlers == {}

    def test_Node_consumer(self):
        consumer1 = self.node.Consumer()
        assert consumer1.channel is self.default_chan
        assert consumer1.no_ack

        chan2 = self.connection.channel()
        consumer2 = self.node.Consumer(channel=chan2, no_ack=False)
        assert consumer2.channel is chan2
        assert not consumer2.no_ack

    def test_Node_consumer_multiple_listeners(self):
        warnings.resetwarnings()
        consumer = self.node.Consumer()
        q = consumer.queues[0]
        with warnings.catch_warnings(record=True) as log:
            q.on_declared('foo', 1, 1)
            assert log
            assert 'already using this' in log[0].message.args[0]

        with warnings.catch_warnings(record=True) as log:
            q.on_declared('foo', 1, 0)
            assert not log

    def test_handler(self):
        node = self.bound.Node('test_handler', state=self.state)

        @node.handler
        def my_handler_name(state):
            return 42

        assert 'my_handler_name' in node.handlers

    def test_dispatch(self):
        node = self.bound.Node('test_dispatch', state=self.state)

        @node.handler
        def my_handler_name(state, x=None, y=None):
            return x + y

        assert node.dispatch('my_handler_name',
                             arguments={'x': 10, 'y': 10}) == 20

    def test_dispatch_raising_SystemExit(self):
        node = self.bound.Node('test_dispatch_raising_SystemExit',
                               state=self.state)

        @node.handler
        def my_handler_name(state):
            raise SystemExit

        with pytest.raises(SystemExit):
            node.dispatch('my_handler_name')

    def test_dispatch_raising(self):
        node = self.bound.Node('test_dispatch_raising', state=self.state)

        @node.handler
        def my_handler_name(state):
            raise KeyError('foo')

        res = node.dispatch('my_handler_name')
        assert 'error' in res
        assert 'KeyError' in res['error']

    def test_dispatch_replies(self):
        _replied = [False]

        def reply(data, **options):
            _replied[0] = True

        node = self.bound.Node('test_dispatch', state=self.state)
        node.reply = reply

        @node.handler
        def my_handler_name(state, x=None, y=None):
            return x + y

        node.dispatch('my_handler_name',
                      arguments={'x': 10, 'y': 10},
                      reply_to={'exchange': 'foo', 'routing_key': 'bar'})
        assert _replied[0]

    def test_reply(self):
        _replied = [(None, None, None)]

        def publish_reply(data, exchange, routing_key, ticket, **kwargs):
            _replied[0] = (data, exchange, routing_key, ticket)

        mailbox = self.mailbox(self.connection)
        mailbox._publish_reply = publish_reply
        node = mailbox.Node('test_reply')

        @node.handler
        def my_handler_name(state):
            return 42

        node.dispatch('my_handler_name',
                      reply_to={'exchange': 'exchange',
                                'routing_key': 'rkey'},
                      ticket='TICKET')
        data, exchange, routing_key, ticket = _replied[0]
        assert data == {'test_reply': 42}
        assert exchange == 'exchange'
        assert routing_key == 'rkey'
        assert ticket == 'TICKET'

    def test_handle_message(self):
        node = self.bound.Node('test_dispatch_from_message')

        @node.handler
        def my_handler_name(state, x=None, y=None):
            return x * y

        body = {'method': 'my_handler_name',
                'arguments': {'x': 64, 'y': 64}}

        assert node.handle_message(body, None) == 64 * 64

        # message not for me should not be processed.
        body['destination'] = ['some_other_node']
        assert node.handle_message(body, None) is None

        # message for me should be processed.
        body['destination'] = ['test_dispatch_from_message']
        assert node.handle_message(body, None) is not None

        # message not for me should not be processed.
        body.pop("destination")
        body['matcher'] = 'glob'
        body["pattern"] = "something*"
        assert node.handle_message(body, None) is None

        body["pattern"] = "test*"
        assert node.handle_message(body, None) is not None

    def test_handle_message_adjusts_clock(self):
        node = self.bound.Node('test_adjusts_clock')

        @node.handler
        def my_handler_name(state):
            return 10

        body = {'method': 'my_handler_name',
                'arguments': {}}
        message = Mock(name='message')
        message.headers = {'clock': 313}
        node.adjust_clock = Mock(name='adjust_clock')
        res = node.handle_message(body, message)
        node.adjust_clock.assert_called_with(313)
        assert res == 10

    def test_listen(self):
        consumer = self.node.listen()
        assert consumer.callbacks[0] == self.node.handle_message
        assert consumer.channel == self.default_chan

    def test_cast(self):
        self.bound.cast(['somenode'], 'mymethod')
        consumer = self.node.Consumer()
        assert is_cast(self.get_next(consumer))

    def test_abcast(self):
        self.bound.abcast('mymethod')
        consumer = self.node.Consumer()
        assert is_cast(self.get_next(consumer))

    def test_call_destination_must_be_sequence(self):
        with pytest.raises(ValueError):
            self.bound.call('some_node', 'mymethod')

    def test_call(self):
        assert self.bound.call(['some_node'], 'mymethod') == 'COLLECTED'
        consumer = self.node.Consumer()
        assert is_call(self.get_next(consumer))

    def test_multi_call(self):
        assert self.bound.multi_call('mymethod') == 'COLLECTED'
        consumer = self.node.Consumer()
        assert is_call(self.get_next(consumer))

    def get_next(self, consumer):
        m = consumer.queues[0].get()
        if m:
            return m.payload


GLOBAL_PIDBOX = None


def getoid():
    return GLOBAL_PIDBOX.oid


class test_PidboxOid:
    """Unittests checking oid consistency of Pidbox"""

    def setup(self):
        global GLOBAL_PIDBOX
        GLOBAL_PIDBOX = pidbox.Mailbox('unittest_mailbox')

    def test_oid_consistency(self):
        """Tests that oid is consistent in single process"""
        m1 = pidbox.Mailbox('mailbox1')
        m2 = pidbox.Mailbox('mailbox2')
        assert m1.oid == m1.oid
        assert m2.oid == m2.oid
        assert m1.oid != m2.oid

    def test_subprocess_oid(self):
        """Tests that subprocess will not share oid with parent process."""
        oid = GLOBAL_PIDBOX.oid
        with ProcessPoolExecutor() as e:
            res = e.submit(getoid)
            subprocess_oid = res.result()
        assert subprocess_oid != oid

    def test_thread_oid(self):
        """Tests that threads will not share oid."""
        oid = GLOBAL_PIDBOX.oid
        with ThreadPoolExecutor() as e:
            res = e.submit(getoid)
            subprocess_oid = res.result()
        assert subprocess_oid != oid