Merge pull request allan-simon#12 from weenect/feature-9

fixes allan-simon#9 - add basic nack support

Author: allan-simon

README.md

@@ -61,6 +61,8 @@ After some searching, it was a bug in the 3rd party library and my code, that wa
 
   * `GET /authentification-done-with-success-on/$USERNAME` return if a client has successfully connected with the given username
   * `GET /messages-acknowledged/$DELIVERY_TAG` this call will block until the given delivery tag was acknowledged as consumed by a worker, if the message has already been consumed before the call, the call directly returns. => this call is super pratical to write your integration tests (i.e to check if you worker has written a record in database after the message was consumed)
+  * `GET /messages-not-acknowledged/$DELIVERY_TAG` this call will block until the given delivery tag was marked as rejected by a worker (basic_nack), if the message has already been rejected before the call, the call directly returns. => this call is super practical to write your integration tests (i.e worker rejected the message for any reason)
+  * `GET /messages-requeued/$DELIVERY_TAG` this call will block until the given delivery tag was marked as rejected by a worker with requeue parameter to true (basic_nack with requeue), if the message has already been requeued before the call, the call directly returns. => this call is super practical to write your integration tests (i.e worker rejected the message to be treated later)
   * `GET /messages-in-queue/$QUEUE_NAME` get all the messages waiting to be consumed in a given queue
   * `GET /messages-in-exchange/$EXCHANGE_NAME` get all the messages waiting to be consumed in a given exchange
   * `GET /queue-bound-to-exchange/$QUEUE_NAME/$EXCHANGE_NAME` wait until a given queue is bound to the given exchange

service/http_protocol.py

@@ -109,6 +109,34 @@ def _on_get(self, target):
             future.add_done_callback(self._on_get_done)
             return
 
+        ###
+        # Wait for a message identified by a delivery_tag to be nack
+        # by the consumer or timeout
+        ###
+        if target.startswith(b'/messages-not-acknowledged/'):
+            delivery_tag = target.split(b'/', maxsplit=2)[2]
+            future = asyncio.ensure_future(
+                self._global_state.wait_message_not_acknowledged(
+                    int(delivery_tag.decode('utf-8')),
+                )
+            )
+            future.add_done_callback(self._on_get_done)
+            return
+
+        ###
+        # Wait for a message identified by a delivery_tag to be nack and requeued
+        # by the consumer or timeout
+        ###
+        if target.startswith(b'/messages-requeued/'):
+            delivery_tag = target.split(b'/', maxsplit=2)[2]
+            future = asyncio.ensure_future(
+                self._global_state.wait_message_requeued(
+                    int(delivery_tag.decode('utf-8')),
+                )
+            )
+            future.add_done_callback(self._on_get_done)
+            return
+
         ###
         # Inspect the content of a queue where the program we test
         # publish messages.
@@ -248,7 +276,6 @@ def _on_put(self, target, data):
         self._send_http_response_not_found()
 
     def _on_get_done(self, future):
-
         try:
             success = future.result()
             if not success:

service/method.py

@@ -18,6 +18,7 @@ class MethodIDs(IntEnum):
     BASIC_PUBLISH = 0x003C0028
     BASIC_CONSUME = 0x003C0014
     BASIC_ACK = 0x003C0050
+    BASIC_NACK = 0x003C0078
 
     EXCHANGE_DECLARE = 0x0028000A
 
@@ -296,6 +297,19 @@ def _decode_basic_ack(payload):
     }
 
 
+def _decode_basic_nack(payload):
+    values, _ = loads(
+        'Lbb',
+        payload,
+        offset=4,
+    )
+    return {
+        'delivery-tag': values[0],
+        'multiple': values[1],
+        'requeue': values[2],
+    }
+
+
 def _decode_basic_cancel(payload):
 
     values, _ = loads(
@@ -323,6 +337,7 @@ def _decode_basic_cancel(payload):
     0x003C0028: _decode_basic_publish,
     MethodIDs.BASIC_CONSUME: _decode_basic_consume,
     MethodIDs.BASIC_ACK: _decode_basic_ack,
+    MethodIDs.BASIC_NACK: _decode_basic_nack,
 
     MethodIDs.BASIC_CANCEL: _decode_basic_cancel,
 

service/protocol.py

@@ -333,6 +333,13 @@ def _treat_channel_frame(self, frame_value):
                 )
                 return
 
+            if frame_value.method_id == MethodIDs.BASIC_NACK:
+                self._global_state.message_nack(
+                    frame_value.properties['delivery-tag'],
+                    frame_value.properties['requeue'],
+                )
+                return
+
             if frame_value.method_id == MethodIDs.BASIC_CANCEL:
                 send_basic_cancel_ok(
                     self.transport,

service/state.py

@@ -22,6 +22,8 @@ def __init__(self):
         self._queues_bound_exhanges = {}
         self._authentication_tried_on = {}
         self._message_acknowledged = set()
+        self._message_not_acknowledged = set()
+        self._message_requeued = set()
 
     def check_credentials(self, username, password):
         is_authentified = self._users.get(username, None) == password
@@ -270,6 +272,12 @@ def publish_message_in_queue(
     def message_ack(self, delivery_tag):
         self._message_acknowledged.add(delivery_tag)
 
+    def message_nack(self, delivery_tag, requeue: bool = False):
+        if requeue:
+            self._message_requeued.add(delivery_tag)
+        else:
+            self._message_not_acknowledged.add(delivery_tag)
+
     async def wait_authentication_performed_on(self, username, timeout=10):
         for _ in range(timeout):
             decoded_username = username.decode('utf-8')
@@ -289,6 +297,24 @@ async def wait_message_acknoledged(self, delivery_tag, timeout=10):
 
         raise WaitTimeout()
 
+    async def wait_message_not_acknowledged(self, delivery_tag, timeout=10):
+        for _ in range(timeout):
+            if delivery_tag in self._message_not_acknowledged:
+                return True
+
+            await asyncio.sleep(1)
+
+        raise WaitTimeout()
+
+    async def wait_message_requeued(self, delivery_tag, timeout=10):
+        for _ in range(timeout):
+            if delivery_tag in self._message_requeued:
+                return True
+
+            await asyncio.sleep(1)
+
+        raise WaitTimeout()
+
     async def wait_queue_bound(self, queue, exchange, timeout=10):
         for _ in range(timeout):