From 4cdcc4d155d1bd11afb6366f2b5efa1739788d01 Mon Sep 17 00:00:00 2001
From: mathieui <mathieui@mathieui.net>
Date: Mon, 8 Mar 2021 22:19:46 +0100
Subject: [PATCH] itests: test mam.iterate() as well for fetching the archive

---
 itests/test_mam.py | 45 +++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 41 insertions(+), 4 deletions(-)

diff --git a/itests/test_mam.py b/itests/test_mam.py
index f61bc1e6..d10241bd 100644
--- a/itests/test_mam.py
+++ b/itests/test_mam.py
@@ -22,11 +22,14 @@ class TestMAM(SlixIntegration):
         """Make sure we can get messages from our archive"""
         # send messages first
         tok = randint(1, 999999)
-        self.clients[0].make_message(mto=self.clients[1].boundjid, mbody='coucou').send()
+        self.clients[0].make_message(
+            mto=self.clients[1].boundjid,
+            mbody=f'coucou {tok}'
+        ).send()
         await self.clients[1].wait_until('message')
         self.clients[1].make_message(
             mto=self.clients[0].boundjid,
-            mbody='coucou coucou %s' % tok,
+            mbody=f'coucou coucou {tok}',
         ).send()
         await self.clients[0].wait_until('message')
 
@@ -48,8 +51,42 @@ class TestMAM(SlixIntegration):
             if count >= 2:
                 break
 
-        self.assertEqual(msgs[0]['body'], 'coucou')
-        self.assertEqual(msgs[1]['body'], 'coucou coucou %s' % tok)
+        self.assertEqual(msgs[0]['body'], f'coucou {tok}')
+        self.assertEqual(msgs[1]['body'], f'coucou coucou {tok}')
+
+    async def test_mam_iterate(self):
+        """Make sure we can iterate over messages from our archive"""
+        # send messages first
+        tok = randint(1, 999999)
+        self.clients[0].make_message(
+            mto=self.clients[1].boundjid,
+            mbody=f'coucou {tok}'
+        ).send()
+        await self.clients[1].wait_until('message')
+        self.clients[1].make_message(
+            mto=self.clients[0].boundjid,
+            mbody='coucou coucou %s' % tok,
+        ).send()
+        await self.clients[0].wait_until('message')
+
+        # Get archive
+        retrieve = self.clients[0]['xep_0313'].iterate(
+            with_jid=JID(self.envjid('CI_ACCOUNT2')),
+            reverse=True,
+            rsm={'max': 1}
+        )
+        msgs = []
+        count = 0
+        async for msg in retrieve:
+            msgs.append(
+                msg['mam_result']['forwarded']['stanza']
+            )
+            count += 1
+            if count >= 2:
+                break
+
+        self.assertEqual(msgs[0]['body'], f'coucou coucou {tok}')
+        self.assertEqual(msgs[1]['body'], f'coucou {tok}')
 
 
 suite = unittest.TestLoader().loadTestsFromTestCase(TestMAM)