Merge pull request #45 from CybercentreCanada/assemblyline/issues/184

Allow paging of messages from notification queues

assemblyline_client/v4_client/module/ingest.py

@@ -114,13 +114,19 @@ def get_message(self, nq):
 """
         return self._connection.get(api_path_by_module(self, nq))
 
-    def get_message_list(self, nq):
+    def get_message_list(self, nq, page_size=None):
         """\
 Return all messages from the given notification queue.
 
 Required:
-nq      : Notification queue name. (string)
+nq          : Notification queue name. (string)
+
+Optional:
+page_size   : Number of messages to get back from queue. (int)
 
 Throws a Client exception if the watch queue does not exist.
 """
-        return self._connection.get(api_path_by_module(self, nq))
+        kw = {}
+        if page_size:
+            kw['page_size'] = int(page_size)
+        return self._connection.get(api_path_by_module(self, nq, **kw))

test/test_ingest.py

@@ -54,6 +54,26 @@ def test_get_message_list(datastore, client):
     assert res[0] == msg_0
     assert res[1] == msg_1
 
+def test_get_message_list_with_paging(datastore, client):
+    notification_queue = get_random_id()
+    queue = NamedQueue("nq-%s" % notification_queue,
+                       host=config.core.redis.persistent.host,
+                       port=config.core.redis.persistent.port)
+    queue.delete()
+    msg_0 = random_model_obj(Submission).as_primitives()
+    queue.push(msg_0)
+    msg_1 = random_model_obj(Submission).as_primitives()
+    queue.push(msg_1)
+
+    res = True
+    messages = []
+    while res:
+        res = client.ingest.get_message_list(notification_queue, page_size=1)
+        messages += res
+
+    assert len(messages) == 2
+    assert messages[0] == msg_0
+    assert messages[1] == msg_1
 
 def test_ingest_content(datastore, client):
     content = get_random_phrase(wmin=15, wmax=50).encode()