EventReceiver: now uses kombu.mixins.ConsumerMixin

celery/events/__init__.py

@@ -19,6 +19,7 @@ from contextlib import contextmanager
 from copy import copy
 
 from kombu import eventloop, Exchange, Queue, Consumer, Producer
+from kombu.mixins import ConsumerMixin
 from kombu.utils import cached_property
 
 from celery.app import app_or_default
@@ -153,7 +154,7 @@ class EventDispatcher(object):
         self.publisher = None
 
 
-class EventReceiver(object):
+class EventReceiver(ConsumerMixin):
     """Capture events.
 
     :param connection: Connection to the broker.
@@ -190,21 +191,17 @@ class EventReceiver(object):
         handler = self.handlers.get(type) or self.handlers.get('*')
         handler and handler(event)
 
-    @contextmanager
-    def consumer(self, wakeup=True):
-        """Create event consumer."""
-        consumer = Consumer(self.connection,
-                            queues=[self.queue], no_ack=True)
-        consumer.register_callback(self._receive)
-        with consumer:
-            if wakeup:
-                self.wakeup_workers(channel=consumer.channel)
-            yield consumer
+    def get_consumers(self, Consumer, channel):
+        return [Consumer(queues=[self.queue],
+                         callbacks=[self._receive], no_ack=True)]
+
+    def on_consume_ready(self, connection, channel, consumers,
+            wakeup=True, **kwargs):
+        if wakeup:
+            self.wakeup_workers(channel=channel)
 
     def itercapture(self, limit=None, timeout=None, wakeup=True):
-        with self.consumer(wakeup=wakeup) as consumer:
-            yield consumer
-            self.drain_events(limit=limit, timeout=timeout)
+        return self.consume(limit=limit, timeout=timeout, wakeup=wakeup)
 
     def capture(self, limit=None, timeout=None, wakeup=True):
         """Open up a consumer capturing events.
@@ -213,17 +210,13 @@ class EventReceiver(object):
         stop unless forced via :exc:`KeyboardInterrupt` or :exc:`SystemExit`.
 
         """
-        list(self.itercapture(limit=limit, timeout=timeout, wakeup=wakeup))
+        return list(self.consume(limit=limit, timeout=timeout, wakeup=wakeup))
 
     def wakeup_workers(self, channel=None):
         self.app.control.broadcast('heartbeat',
                                    connection=self.connection,
                                    channel=channel)
 
-    def drain_events(self, **kwargs):
-        for _ in eventloop(self.connection, **kwargs):
-            pass
-
     def _receive(self, body, message):
         type = body.pop('type').lower()
         clock = body.get('clock')

celery/tests/events/test_events.py

@@ -163,9 +163,6 @@ class test_EventReceiver(AppCase):
         try:
             r = self.app.events.Receiver(connection, node_id='celery.tests')
             it = r.itercapture(timeout=0.0001, wakeup=False)
-            consumer = next(it)
-            self.assertTrue(consumer.queues)
-            self.assertEqual(consumer.callbacks[0], r._receive)
 
             with self.assertRaises(socket.timeout):
                 next(it)