Merge pull request #49 from smagafurov/feature/message-process-contex…

…t-check

ability to process message yourself under message.process context

aio_pika/message.py

@@ -297,7 +297,7 @@ def __init__(self, channel: Channel, envelope, properties, body, no_ack: bool =
             self.__processed = True
 
     @contextmanager
-    def process(self, requeue=False, reject_on_redelivered=False):
+    def process(self, requeue=False, reject_on_redelivered=False, ignore_processed=False):
         """ Context manager for processing the message
 
             >>> def on_message_received(message: IncomingMessage):
@@ -306,19 +306,31 @@ def process(self, requeue=False, reject_on_redelivered=False):
             ...        # the message will be rejected
             ...        print(message.body)
 
+        Example with ignore_processed=True
+
+            >>> def on_message_received(message: IncomingMessage):
+            ...    with message.process(ignore_processed=True):
+            ...        # Now (with ignore_processed=True) you may reject (or ack) message manually too
+            ...        if True:  # some reasonable condition here
+            ...            message.reject()
+            ...        print(message.body)
+
         :param requeue: Requeue message when exception.
         :param reject_on_redelivered: When True message will be rejected only when message was redelivered.
+        :param ignore_processed: Do nothing if message already processed
 
         """
         try:
             yield self
-            self.ack()
+            if not ignore_processed or not self.processed:
+                self.ack()
         except:
-            if reject_on_redelivered and self.redelivered:
-                log.info("Message %r was redelivered and will be rejected.", self)
-                self.reject(requeue=False)
-            else:
-                self.reject(requeue=requeue)
+            if not ignore_processed or not self.processed:
+                if reject_on_redelivered and self.redelivered:
+                    log.info("Message %r was redelivered and will be rejected.", self)
+                    self.reject(requeue=False)
+                else:
+                    self.reject(requeue=requeue)
             raise
 
     def ack(self, multiple: bool = False):

tests/test_amqp.py

@@ -378,6 +378,29 @@ def test_context_process(self):
             routing_key
         )
 
+        incoming_message = yield from queue.get(timeout=5)
+
+        with self.assertRaises(MessageProcessError):
+            with incoming_message.process():
+                incoming_message.reject(requeue=True)
+
+        self.assertEqual(incoming_message.locked, True)
+
+        incoming_message = yield from queue.get(timeout=5)
+
+        with incoming_message.process(ignore_processed=True):
+            incoming_message.reject(requeue=False)
+
+        self.assertEqual(incoming_message.body, body)
+
+        yield from exchange.publish(
+            Message(
+                body, content_type='text/plain',
+                headers={'foo': 'bar'}
+            ),
+            routing_key
+        )
+
         incoming_message = yield from queue.get(timeout=5)
         with self.assertRaises(AssertionError):
             with incoming_message.process(requeue=True, reject_on_redelivered=True):
@@ -859,7 +882,7 @@ def test_connection_close(self):
 
         routing_key = self.get_random_name()
 
-        channel = yield from client.channel()    # type: Channel
+        channel = yield from client.channel()    # type: aio_pika.Channel
         exchange = yield from channel.declare_exchange('direct', auto_delete=True)
 
         try: