import errno
import logging
import socket
from collections import deque
from unittest.mock import MagicMock, Mock, call, patch

import pytest
from amqp import ChannelError
from billiard.exceptions import RestartFreqExceeded

from celery import bootsteps
from celery.contrib.testing.mocks import ContextMock
from celery.exceptions import WorkerShutdown, WorkerTerminate
from celery.utils.collections import LimitedSet
from celery.utils.quorum_queues import detect_quorum_queues
from celery.worker.consumer.agent import Agent
from celery.worker.consumer.consumer import CANCEL_TASKS_BY_DEFAULT, CLOSE, TERMINATE, Consumer
from celery.worker.consumer.gossip import Gossip
from celery.worker.consumer.heart import Heart
from celery.worker.consumer.mingle import Mingle
from celery.worker.consumer.tasks import Tasks
from celery.worker.state import active_requests


class ConsumerTestCase:
    def get_consumer(self, no_hub=False, **kwargs):
        consumer = Consumer(
            on_task_request=Mock(),
            init_callback=Mock(),
            pool=Mock(),
            app=self.app,
            timer=Mock(),
            controller=Mock(),
            hub=None if no_hub else Mock(),
            **kwargs
        )
        consumer.blueprint = Mock(name='blueprint')
        consumer.pool.num_processes = 2
        consumer._restart_state = Mock(name='_restart_state')
        consumer.connection = _amqp_connection()
        consumer.connection_errors = (socket.error, OSError,)
        consumer.conninfo = consumer.connection
        return consumer


class test_Consumer(ConsumerTestCase):
    def setup_method(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        self.add = add

    def test_repr(self):
        assert repr(self.get_consumer())

    def test_taskbuckets_defaultdict(self):
        c = self.get_consumer()
        assert c.task_buckets['fooxasdwx.wewe'] is None

    def test_sets_heartbeat(self):
        c = self.get_consumer(amqheartbeat=10)
        assert c.amqheartbeat == 10
        self.app.conf.broker_heartbeat = 20
        c = self.get_consumer(amqheartbeat=None)
        assert c.amqheartbeat == 20

    def test_gevent_bug_disables_connection_timeout(self):
        with patch('celery.worker.consumer.consumer._detect_environment') as d:
            d.return_value = 'gevent'
            self.app.conf.broker_connection_timeout = 33.33
            self.get_consumer()
            assert self.app.conf.broker_connection_timeout is None

    def test_limit_moved_to_pool(self):
        with patch('celery.worker.consumer.consumer.task_reserved') as task_reserved:
            c = self.get_consumer()
            c.on_task_request = Mock(name='on_task_request')
            request = Mock(name='request')
            c._limit_move_to_pool(request)
            task_reserved.assert_called_with(request)
            c.on_task_request.assert_called_with(request)

    def test_update_prefetch_count(self):
        c = self.get_consumer()
        c._update_qos_eventually = Mock(name='update_qos')
        c.initial_prefetch_count = None
        c.pool.num_processes = None
        c.prefetch_multiplier = 10
        assert c._update_prefetch_count(1) is None
        c.initial_prefetch_count = 10
        c.pool.num_processes = 10
        c._update_prefetch_count(8)
        c._update_qos_eventually.assert_called_with(8)
        assert c.initial_prefetch_count == 10 * 10

    @pytest.mark.parametrize(
        'active_requests_count,expected_initial,expected_maximum,enabled',
        [
            [0, 2, True, True],
            [1, 1, False, True],
            [2, 1, False, True],
            [0, 2, True, False],
            [1, 2, True, False],
            [2, 2, True, False],
        ]
    )
    @patch('celery.worker.consumer.consumer.active_requests', new_callable=set)
    def test_restore_prefetch_count_on_restart(self, active_requests_mock, active_requests_count,
                                               expected_initial, expected_maximum, enabled, subtests):
        self.app.conf.worker_enable_prefetch_count_reduction = enabled

        reqs = {Mock() for _ in range(active_requests_count)}
        active_requests_mock.update(reqs)

        c = self.get_consumer()
        c.qos = Mock()
        c.blueprint = Mock()

        def bp_start(*_, **__):
            if c.restart_count > 1:
                c.blueprint.state = CLOSE
            else:
                raise ConnectionError

        c.blueprint.start.side_effect = bp_start

        c.start()

        with subtests.test("initial prefetch count is never 0"):
            assert c.initial_prefetch_count != 0

        with subtests.test(f"initial prefetch count is equal to {expected_initial}"):
            assert c.initial_prefetch_count == expected_initial

        with subtests.test("maximum prefetch is reached"):
            assert c._maximum_prefetch_restored is expected_maximum

    def test_restore_prefetch_count_after_connection_restart_negative(self):
        self.app.conf.worker_enable_prefetch_count_reduction = False

        c = self.get_consumer()
        c.qos = Mock()

        # Overcome TypeError: 'Mock' object does not support the context manager protocol
        class MutexMock:
            def __enter__(self):
                pass

            def __exit__(self, *args):
                pass

        c.qos._mutex = MutexMock()

        assert c._restore_prefetch_count_after_connection_restart(None) is None

    def test_create_task_handler(self, subtests):
        c = self.get_consumer()
        c.qos = MagicMock()
        c.qos.value = 1
        c._maximum_prefetch_restored = False

        sig = self.add.s(2, 2)
        message = self.task_message_from_sig(self.app, sig)

        def raise_exception():
            raise KeyError('Foo')

        def strategy(_, __, ack_log_error_promise, ___, ____):
            ack_log_error_promise()

        c.strategies[sig.task] = strategy
        c.call_soon = raise_exception
        on_task_received = c.create_task_handler()
        on_task_received(message)

        with subtests.test("initial prefetch count is never 0"):
            assert c.initial_prefetch_count != 0

        with subtests.test("initial prefetch count is 2"):
            assert c.initial_prefetch_count == 2

        with subtests.test("maximum prefetch is reached"):
            assert c._maximum_prefetch_restored is True

    def test_flush_events(self):
        c = self.get_consumer()
        c.event_dispatcher = None
        c._flush_events()
        c.event_dispatcher = Mock(name='evd')
        c._flush_events()
        c.event_dispatcher.flush.assert_called_with()

    def test_on_send_event_buffered(self):
        c = self.get_consumer()
        c.hub = None
        c.on_send_event_buffered()
        c.hub = Mock(name='hub')
        c.on_send_event_buffered()
        c.hub._ready.add.assert_called_with(c._flush_events)

    def test_schedule_bucket_request(self):
        c = self.get_consumer()
        c.timer = Mock()

        bucket = Mock()
        request = Mock()
        bucket.pop = lambda: bucket.contents.popleft()
        bucket.can_consume.return_value = True
        bucket.contents = deque()

        with patch(
            'celery.worker.consumer.consumer.Consumer._limit_move_to_pool'
        ) as task_reserved:
            bucket.contents.append((request, 3))
            c._schedule_bucket_request(bucket)
            bucket.can_consume.assert_called_with(3)
            task_reserved.assert_called_with(request)

        bucket.can_consume.return_value = False
        bucket.contents = deque()
        bucket.expected_time.return_value = 3.33
        bucket.contents.append((request, 4))
        limit_order = c._limit_order
        c._schedule_bucket_request(bucket)
        assert c._limit_order == limit_order + 1
        bucket.can_consume.assert_called_with(4)
        c.timer.call_after.assert_called_with(
            3.33, c._schedule_bucket_request, (bucket,),
            priority=c._limit_order,
        )
        bucket.expected_time.assert_called_with(4)
        assert bucket.pop() == (request, 4)

        bucket.contents = deque()
        bucket.can_consume.reset_mock()
        c._schedule_bucket_request(bucket)
        bucket.can_consume.assert_not_called()

    def test_limit_task(self):
        c = self.get_consumer()
        bucket = Mock()
        request = Mock()

        with patch(
            'celery.worker.consumer.consumer.Consumer._schedule_bucket_request'
        ) as task_reserved:
            c._limit_task(request, bucket, 1)
            bucket.add.assert_called_with((request, 1))
            task_reserved.assert_called_with(bucket)

    def test_post_eta(self):
        c = self.get_consumer()
        c.qos = Mock()
        bucket = Mock()
        request = Mock()

        with patch(
            'celery.worker.consumer.consumer.Consumer._schedule_bucket_request'
        ) as task_reserved:
            c._limit_post_eta(request, bucket, 1)
            c.qos.decrement_eventually.assert_called_with()
            bucket.add.assert_called_with((request, 1))
            task_reserved.assert_called_with(bucket)

    def test_max_restarts_exceeded(self):
        c = self.get_consumer()

        def se(*args, **kwargs):
            c.blueprint.state = CLOSE
            raise RestartFreqExceeded()

        c._restart_state.step.side_effect = se
        c.blueprint.start.side_effect = socket.error()

        with patch('celery.worker.consumer.consumer.sleep') as sleep:
            c.start()
            sleep.assert_called_with(1)

    def test_do_not_restart_when_closed(self):
        c = self.get_consumer()

        c.blueprint.state = None

        def bp_start(*args, **kwargs):
            c.blueprint.state = CLOSE

        c.blueprint.start.side_effect = bp_start
        with patch('celery.worker.consumer.consumer.sleep'):
            c.start()

        c.blueprint.start.assert_called_once_with(c)

    def test_do_not_restart_when_terminated(self):
        c = self.get_consumer()

        c.blueprint.state = None

        def bp_start(*args, **kwargs):
            c.blueprint.state = TERMINATE

        c.blueprint.start.side_effect = bp_start

        with patch('celery.worker.consumer.consumer.sleep'):
            c.start()

        c.blueprint.start.assert_called_once_with(c)

    def test_too_many_open_files_raises_error(self):
        c = self.get_consumer()
        err = OSError()
        err.errno = errno.EMFILE
        c.blueprint.start.side_effect = err
        with pytest.raises(WorkerTerminate):
            c.start()

    def _closer(self, c):
        def se(*args, **kwargs):
            c.blueprint.state = CLOSE

        return se

    @pytest.mark.parametrize("broker_connection_retry", [True, False])
    def test_blueprint_restart_when_state_not_in_stop_conditions(self, broker_connection_retry):
        c = self.get_consumer()

        # ensure that WorkerShutdown is not raised
        c.app.conf['broker_connection_retry'] = broker_connection_retry
        c.app.conf['broker_connection_retry_on_startup'] = True
        c.restart_count = -1

        # ensure that blueprint state is not in stop conditions
        c.blueprint.state = bootsteps.RUN
        c.blueprint.start.side_effect = ConnectionError()

        # stops test from running indefinitely in the while loop
        c.blueprint.restart.side_effect = self._closer(c)

        c.start()
        c.blueprint.restart.assert_called_once()

    @pytest.mark.parametrize("broker_channel_error_retry", [True, False])
    def test_blueprint_restart_for_channel_errors(self, broker_channel_error_retry):
        c = self.get_consumer()

        # ensure that WorkerShutdown is not raised
        c.app.conf['broker_connection_retry'] = True
        c.app.conf['broker_connection_retry_on_startup'] = True
        c.app.conf['broker_channel_error_retry'] = broker_channel_error_retry
        c.restart_count = -1

        # ensure that blueprint state is not in stop conditions
        c.blueprint.state = bootsteps.RUN
        c.blueprint.start.side_effect = ChannelError()

        # stops test from running indefinitely in the while loop
        c.blueprint.restart.side_effect = self._closer(c)

        # restarted only when broker_channel_error_retry is True
        if broker_channel_error_retry:
            c.start()
            c.blueprint.restart.assert_called_once()
        else:
            with pytest.raises(ChannelError):
                c.start()

    def test_collects_at_restart(self):
        c = self.get_consumer()
        c.connection.collect.side_effect = MemoryError()
        c.blueprint.start.side_effect = socket.error()
        c.blueprint.restart.side_effect = self._closer(c)
        c.start()
        c.connection.collect.assert_called_with()

    def test_register_with_event_loop(self):
        c = self.get_consumer()
        c.register_with_event_loop(Mock(name='loop'))

    def test_on_close_clears_semaphore_timer_and_reqs(self):
        with patch('celery.worker.consumer.consumer.reserved_requests') as res:
            c = self.get_consumer()
            c.on_close()
            c.controller.semaphore.clear.assert_called_with()
            c.timer.clear.assert_called_with()
            res.clear.assert_called_with()
            c.pool.flush.assert_called_with()

            c.controller = None
            c.timer = None
            c.pool = None
            c.on_close()

    def test_connect_error_handler(self):
        self.app._connection = _amqp_connection()
        conn = self.app._connection.return_value
        c = self.get_consumer()
        assert c.connect()
        conn.ensure_connection.assert_called()
        errback = conn.ensure_connection.call_args[0][0]
        errback(Mock(), 0)

    @patch('celery.worker.consumer.consumer.error')
    def test_connect_error_handler_progress(self, error):
        self.app.conf.broker_connection_retry = True
        self.app.conf.broker_connection_max_retries = 3
        self.app._connection = _amqp_connection()
        conn = self.app._connection.return_value
        c = self.get_consumer()
        assert c.connect()
        errback = conn.ensure_connection.call_args[0][0]
        errback(Mock(), 2)
        assert error.call_args[0][3] == 'Trying again in 2.00 seconds... (1/3)'
        errback(Mock(), 4)
        assert error.call_args[0][3] == 'Trying again in 4.00 seconds... (2/3)'
        errback(Mock(), 6)
        assert error.call_args[0][3] == 'Trying again in 6.00 seconds... (3/3)'

    def test_cancel_long_running_tasks_on_connection_loss(self):
        c = self.get_consumer()
        c.app.conf.worker_cancel_long_running_tasks_on_connection_loss = True

        mock_request_acks_late_not_acknowledged = Mock()
        mock_request_acks_late_not_acknowledged.task.acks_late = True
        mock_request_acks_late_not_acknowledged.acknowledged = False
        mock_request_acks_late_acknowledged = Mock()
        mock_request_acks_late_acknowledged.task.acks_late = True
        mock_request_acks_late_acknowledged.acknowledged = True
        mock_request_acks_early = Mock()
        mock_request_acks_early.task.acks_late = False
        mock_request_acks_early.acknowledged = False

        active_requests.add(mock_request_acks_late_not_acknowledged)
        active_requests.add(mock_request_acks_late_acknowledged)
        active_requests.add(mock_request_acks_early)

        c.on_connection_error_after_connected(Mock())

        mock_request_acks_late_not_acknowledged.cancel.assert_called_once_with(c.pool)
        mock_request_acks_late_acknowledged.cancel.assert_not_called()
        mock_request_acks_early.cancel.assert_not_called()

        active_requests.clear()

    def test_cancel_long_running_tasks_on_connection_loss__warning(self):
        c = self.get_consumer()
        c.app.conf.worker_cancel_long_running_tasks_on_connection_loss = False

        with pytest.deprecated_call(match=CANCEL_TASKS_BY_DEFAULT):
            c.on_connection_error_after_connected(Mock())

    @pytest.mark.usefixtures('depends_on_current_app')
    def test_cancel_all_unacked_requests(self):
        c = self.get_consumer()

        mock_request_acks_late_not_acknowledged = Mock(id='1')
        mock_request_acks_late_not_acknowledged.task.acks_late = True
        mock_request_acks_late_not_acknowledged.acknowledged = False
        mock_request_acks_late_acknowledged = Mock(id='2')
        mock_request_acks_late_acknowledged.task.acks_late = True
        mock_request_acks_late_acknowledged.acknowledged = True
        mock_request_acks_early = Mock(id='3')
        mock_request_acks_early.task.acks_late = False

        active_requests.add(mock_request_acks_late_not_acknowledged)
        active_requests.add(mock_request_acks_late_acknowledged)
        active_requests.add(mock_request_acks_early)

        c.cancel_all_unacked_requests()

        mock_request_acks_late_not_acknowledged.cancel.assert_called_once_with(c.pool)
        mock_request_acks_late_acknowledged.cancel.assert_not_called()
        mock_request_acks_early.cancel.assert_called_once_with(c.pool)

        active_requests.clear()

    @pytest.mark.parametrize("broker_connection_retry", [True, False])
    @pytest.mark.parametrize("broker_connection_retry_on_startup", [None, False])
    @pytest.mark.parametrize("first_connection_attempt", [True, False])
    def test_ensure_connected(self, subtests, broker_connection_retry, broker_connection_retry_on_startup,
                              first_connection_attempt):
        c = self.get_consumer()
        c.first_connection_attempt = first_connection_attempt
        c.app.conf.broker_connection_retry_on_startup = broker_connection_retry_on_startup
        c.app.conf.broker_connection_retry = broker_connection_retry

        if broker_connection_retry is False:
            if broker_connection_retry_on_startup is None:
                with subtests.test("Deprecation warning when startup is None"):
                    with pytest.deprecated_call():
                        c.ensure_connected(Mock())

            with subtests.test("Does not retry when connect throws an error and retry is set to false"):
                conn = Mock()
                conn.connect.side_effect = ConnectionError()
                with pytest.raises(ConnectionError):
                    c.ensure_connected(conn)


@pytest.mark.parametrize(
    "broker_connection_retry_on_startup,is_connection_loss_on_startup",
    [
        pytest.param(False, True, id='shutdown on connection loss on startup'),
        pytest.param(None, True, id='shutdown on connection loss on startup when retry on startup is undefined'),
        pytest.param(False, False, id='shutdown on connection loss not on startup but startup is defined as false'),
        pytest.param(None, False, id='shutdown on connection loss not on startup and startup is not defined'),
        pytest.param(True, False, id='shutdown on connection loss not on startup but startup is defined as true'),
    ]
)
class test_Consumer_WorkerShutdown(ConsumerTestCase):

    def test_start_raises_connection_error(self,
                                           broker_connection_retry_on_startup,
                                           is_connection_loss_on_startup,
                                           caplog, subtests):
        c = self.get_consumer()
        c.first_connection_attempt = True if is_connection_loss_on_startup else False
        c.app.conf['broker_connection_retry'] = False
        c.app.conf['broker_connection_retry_on_startup'] = broker_connection_retry_on_startup
        c.blueprint.start.side_effect = ConnectionError()

        with subtests.test("Consumer raises WorkerShutdown on connection restart"):
            with pytest.raises(WorkerShutdown):
                c.start()

        record = caplog.records[0]
        with subtests.test("Critical error log message is outputted to the screen"):
            assert record.levelname == "CRITICAL"
            action = "establish" if is_connection_loss_on_startup else "re-establish"
            expected_prefix = f"Retrying to {action}"
            assert record.msg.startswith(expected_prefix)
            conn_type_name = c._get_connection_retry_type(
                is_connection_loss_on_startup
            )
            expected_connection_retry_type = f"app.conf.{conn_type_name}=False"
            assert expected_connection_retry_type in record.msg


class test_Consumer_PerformPendingOperations(ConsumerTestCase):

    def test_perform_pending_operations_all_success(self):
        """
        Test that all pending operations are processed successfully when `once=False`.
        """
        c = self.get_consumer(no_hub=True)

        # Create mock operations
        mock_operation_1 = Mock()
        mock_operation_2 = Mock()

        # Add mock operations to _pending_operations
        c._pending_operations = [mock_operation_1, mock_operation_2]

        # Call perform_pending_operations
        c.perform_pending_operations()

        # Assert that all operations were called
        mock_operation_1.assert_called_once()
        mock_operation_2.assert_called_once()

        # Ensure all pending operations are cleared
        assert len(c._pending_operations) == 0

    def test_perform_pending_operations_with_exception(self):
        """
        Test that pending operations are processed even if one raises an exception, and
        the exception is logged when `once=False`.
        """
        c = self.get_consumer(no_hub=True)

        # Mock operations: one failing, one successful
        mock_operation_fail = Mock(side_effect=Exception("Test Exception"))
        mock_operation_success = Mock()

        # Add operations to _pending_operations
        c._pending_operations = [mock_operation_fail, mock_operation_success]

        # Patch logger to avoid logging during the test
        with patch('celery.worker.consumer.consumer.logger.exception') as mock_logger:
            # Call perform_pending_operations
            c.perform_pending_operations()

            # Assert that both operations were attempted
            mock_operation_fail.assert_called_once()
            mock_operation_success.assert_called_once()

            # Ensure the exception was logged
            mock_logger.assert_called_once()

            # Ensure all pending operations are cleared
            assert len(c._pending_operations) == 0


class test_Heart:

    def test_start(self):
        c = Mock()
        c.timer = Mock()
        c.event_dispatcher = Mock()

        with patch('celery.worker.heartbeat.Heart') as hcls:
            h = Heart(c)
            assert h.enabled
            assert h.heartbeat_interval is None
            assert c.heart is None

            h.start(c)
            assert c.heart
            hcls.assert_called_with(c.timer, c.event_dispatcher,
                                    h.heartbeat_interval)
            c.heart.start.assert_called_with()

    def test_start_heartbeat_interval(self):
        c = Mock()
        c.timer = Mock()
        c.event_dispatcher = Mock()

        with patch('celery.worker.heartbeat.Heart') as hcls:
            h = Heart(c, False, 20)
            assert h.enabled
            assert h.heartbeat_interval == 20
            assert c.heart is None

            h.start(c)
            assert c.heart
            hcls.assert_called_with(c.timer, c.event_dispatcher,
                                    h.heartbeat_interval)
            c.heart.start.assert_called_with()


class test_Tasks:

    def setup_method(self):
        self.c = Mock()
        self.c.app.conf.worker_detect_quorum_queues = True
        self.c.connection.qos_semantics_matches_spec = False

    def test_stop(self):
        c = self.c
        tasks = Tasks(c)
        assert c.task_consumer is None
        assert c.qos is None

        c.task_consumer = Mock()
        tasks.stop(c)

    def test_stop_already_stopped(self):
        c = self.c
        tasks = Tasks(c)
        tasks.stop(c)

    def test_detect_quorum_queues_positive(self):
        c = self.c
        self.c.connection.transport.driver_type = 'amqp'
        c.app.amqp.queues = {"celery": Mock(queue_arguments={"x-queue-type": "quorum"})}
        result, name = detect_quorum_queues(c.app, c.connection.transport.driver_type)
        assert result
        assert name == "celery"

    def test_detect_quorum_queues_negative(self):
        c = self.c
        self.c.connection.transport.driver_type = 'amqp'
        c.app.amqp.queues = {"celery": Mock(queue_arguments=None)}
        result, name = detect_quorum_queues(c.app, c.connection.transport.driver_type)
        assert not result
        assert name == ""

    def test_detect_quorum_queues_not_rabbitmq(self):
        c = self.c
        self.c.connection.transport.driver_type = 'redis'
        result, name = detect_quorum_queues(c.app, c.connection.transport.driver_type)
        assert not result
        assert name == ""

    def test_qos_global_worker_detect_quorum_queues_false(self):
        c = self.c
        c.app.conf.worker_detect_quorum_queues = False
        tasks = Tasks(c)
        assert tasks.qos_global(c) is True

    def test_qos_global_worker_detect_quorum_queues_true_no_quorum_queues(self):
        c = self.c
        c.app.amqp.queues = {"celery": Mock(queue_arguments=None)}
        tasks = Tasks(c)
        assert tasks.qos_global(c) is True

    def test_qos_global_worker_detect_quorum_queues_true_with_quorum_queues(self):
        c = self.c
        self.c.connection.transport.driver_type = 'amqp'
        c.app.amqp.queues = {"celery": Mock(queue_arguments={"x-queue-type": "quorum"})}
        tasks = Tasks(c)
        assert tasks.qos_global(c) is False

    def test_log_when_qos_is_false(self, caplog):
        c = self.c
        c.connection.transport.driver_type = 'amqp'
        c.app.conf.broker_native_delayed_delivery = True
        c.app.amqp.queues = {"celery": Mock(queue_arguments={"x-queue-type": "quorum"})}
        tasks = Tasks(c)

        with caplog.at_level(logging.INFO):
            tasks.start(c)

        assert len(caplog.records) == 1

        record = caplog.records[0]
        assert record.levelname == "INFO"
        assert record.msg == "Global QoS is disabled. Prefetch count in now static."


class test_Agent:

    def test_start(self):
        c = Mock()
        agent = Agent(c)
        agent.instantiate = Mock()
        agent.agent_cls = 'foo:Agent'
        assert agent.create(c) is not None
        agent.instantiate.assert_called_with(agent.agent_cls, c.connection)


class test_Mingle:

    def test_start_no_replies(self):
        c = Mock()
        c.app.connection_for_read = _amqp_connection()
        mingle = Mingle(c)
        I = c.app.control.inspect.return_value = Mock()
        I.hello.return_value = {}
        mingle.start(c)

    def test_start(self):
        c = Mock()
        c.app.connection_for_read = _amqp_connection()
        mingle = Mingle(c)
        assert mingle.enabled

        Aig = LimitedSet()
        Big = LimitedSet()
        Aig.add('Aig-1')
        Aig.add('Aig-2')
        Big.add('Big-1')

        I = c.app.control.inspect.return_value = Mock()
        I.hello.return_value = {
            'A@example.com': {
                'clock': 312,
                'revoked': Aig._data,
            },
            'B@example.com': {
                'clock': 29,
                'revoked': Big._data,
            },
            'C@example.com': {
                'error': 'unknown method',
            },
        }

        our_revoked = c.controller.state.revoked = LimitedSet()

        mingle.start(c)
        I.hello.assert_called_with(c.hostname, our_revoked._data)
        c.app.clock.adjust.assert_has_calls([
            call(312), call(29),
        ], any_order=True)
        assert 'Aig-1' in our_revoked
        assert 'Aig-2' in our_revoked
        assert 'Big-1' in our_revoked


def _amqp_connection():
    connection = ContextMock(name='Connection')
    connection.return_value = ContextMock(name='connection')
    connection.return_value.transport.driver_type = 'amqp'
    return connection


class test_Gossip:

    def test_init(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        assert g.enabled
        assert c.gossip is g

    def test_callbacks(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        on_node_join = Mock(name='on_node_join')
        on_node_join2 = Mock(name='on_node_join2')
        on_node_leave = Mock(name='on_node_leave')
        on_node_lost = Mock(name='on.node_lost')
        g.on.node_join.add(on_node_join)
        g.on.node_join.add(on_node_join2)
        g.on.node_leave.add(on_node_leave)
        g.on.node_lost.add(on_node_lost)

        worker = Mock(name='worker')
        g.on_node_join(worker)
        on_node_join.assert_called_with(worker)
        on_node_join2.assert_called_with(worker)
        g.on_node_leave(worker)
        on_node_leave.assert_called_with(worker)
        g.on_node_lost(worker)
        on_node_lost.assert_called_with(worker)

    def test_election(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        g.start(c)
        g.election('id', 'topic', 'action')
        assert g.consensus_replies['id'] == []
        g.dispatcher.send.assert_called_with(
            'worker-elect', id='id', topic='topic', cver=1, action='action',
        )

    def test_call_task(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        g.start(c)
        signature = g.app.signature = Mock(name='app.signature')
        task = Mock()
        g.call_task(task)
        signature.assert_called_with(task)
        signature.return_value.apply_async.assert_called_with()

        signature.return_value.apply_async.side_effect = MemoryError()
        with patch('celery.worker.consumer.gossip.logger') as logger:
            g.call_task(task)
            logger.exception.assert_called()

    def Event(self, id='id', clock=312,
              hostname='foo@example.com', pid=4312,
              topic='topic', action='action', cver=1):
        return {
            'id': id,
            'clock': clock,
            'hostname': hostname,
            'pid': pid,
            'topic': topic,
            'action': action,
            'cver': cver,
        }

    def test_on_elect(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        g.start(c)

        event = self.Event('id1')
        g.on_elect(event)
        in_heap = g.consensus_requests['id1']
        assert in_heap
        g.dispatcher.send.assert_called_with('worker-elect-ack', id='id1')

        event.pop('clock')
        with patch('celery.worker.consumer.gossip.logger') as logger:
            g.on_elect(event)
            logger.exception.assert_called()

    def Consumer(self, hostname='foo@x.com', pid=4312):
        c = Mock()
        c.app.connection = _amqp_connection()
        c.hostname = hostname
        c.pid = pid
        c.app.events.Receiver.return_value = Mock(accept=[])
        return c

    def setup_election(self, g, c):
        g.start(c)
        g.clock = self.app.clock
        assert 'idx' not in g.consensus_replies
        assert g.on_elect_ack({'id': 'idx'}) is None

        g.state.alive_workers.return_value = [
            'foo@x.com', 'bar@x.com', 'baz@x.com',
        ]
        g.consensus_replies['id1'] = []
        g.consensus_requests['id1'] = []
        e1 = self.Event('id1', 1, 'foo@x.com')
        e2 = self.Event('id1', 2, 'bar@x.com')
        e3 = self.Event('id1', 3, 'baz@x.com')
        g.on_elect(e1)
        g.on_elect(e2)
        g.on_elect(e3)
        assert len(g.consensus_requests['id1']) == 3

        with patch('celery.worker.consumer.gossip.info'):
            g.on_elect_ack(e1)
            assert len(g.consensus_replies['id1']) == 1
            g.on_elect_ack(e2)
            assert len(g.consensus_replies['id1']) == 2
            g.on_elect_ack(e3)
            with pytest.raises(KeyError):
                g.consensus_replies['id1']

    def test_on_elect_ack_win(self):
        c = self.Consumer(hostname='foo@x.com')  # I will win
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        handler = g.election_handlers['topic'] = Mock()
        self.setup_election(g, c)
        handler.assert_called_with('action')

    def test_on_elect_ack_lose(self):
        c = self.Consumer(hostname='bar@x.com')  # I will lose
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        handler = g.election_handlers['topic'] = Mock()
        self.setup_election(g, c)
        handler.assert_not_called()

    def test_on_elect_ack_win_but_no_action(self):
        c = self.Consumer(hostname='foo@x.com')  # I will win
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        g.election_handlers = {}
        with patch('celery.worker.consumer.gossip.logger') as logger:
            self.setup_election(g, c)
            logger.exception.assert_called()

    def test_on_node_join(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        with patch('celery.worker.consumer.gossip.debug') as debug:
            g.on_node_join(c)
            debug.assert_called_with('%s joined the party', 'foo@x.com')

    def test_on_node_leave(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        with patch('celery.worker.consumer.gossip.debug') as debug:
            g.on_node_leave(c)
            debug.assert_called_with('%s left', 'foo@x.com')

    def test_on_node_lost(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        with patch('celery.worker.consumer.gossip.info') as info:
            g.on_node_lost(c)
            info.assert_called_with('missed heartbeat from %s', 'foo@x.com')

    def test_register_timer(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        g.register_timer()
        c.timer.call_repeatedly.assert_called_with(g.interval, g.periodic)
        tref = g._tref
        g.register_timer()
        tref.cancel.assert_called_with()

    def test_periodic(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        g.on_node_lost = Mock()
        state = g.state = Mock()
        worker = Mock()
        state.workers = {'foo': worker}
        worker.alive = True
        worker.hostname = 'foo'
        g.periodic()

        worker.alive = False
        g.periodic()
        g.on_node_lost.assert_called_with(worker)
        with pytest.raises(KeyError):
            state.workers['foo']

    def test_on_message__task(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        assert g.enabled
        message = Mock(name='message')
        message.delivery_info = {'routing_key': 'task.failed'}
        g.on_message(Mock(name='prepare'), message)

    def test_on_message(self):
        c = self.Consumer()
        c.app.connection_for_read = _amqp_connection()
        g = Gossip(c)
        assert g.enabled
        prepare = Mock()
        prepare.return_value = 'worker-online', {}
        c.app.events.State.assert_called_with(
            on_node_join=g.on_node_join,
            on_node_leave=g.on_node_leave,
            max_tasks_in_memory=1,
        )
        g.update_state = Mock()
        worker = Mock()
        g.on_node_join = Mock()
        g.on_node_leave = Mock()
        g.update_state.return_value = worker, 1
        message = Mock()
        message.delivery_info = {'routing_key': 'worker-online'}
        message.headers = {'hostname': 'other'}

        handler = g.event_handlers['worker-online'] = Mock()
        g.on_message(prepare, message)
        handler.assert_called_with(message.payload)
        g.event_handlers = {}

        g.on_message(prepare, message)

        message.delivery_info = {'routing_key': 'worker-offline'}
        prepare.return_value = 'worker-offline', {}
        g.on_message(prepare, message)

        message.delivery_info = {'routing_key': 'worker-baz'}
        prepare.return_value = 'worker-baz', {}
        g.update_state.return_value = worker, 0
        g.on_message(prepare, message)

        message.headers = {'hostname': g.hostname}
        g.on_message(prepare, message)
        g.clock.forward.assert_called_with()