Merge pull request #111 from PeterJCLaw/fix-flake

Fix a couple of flaky tests which depended on partition ordering
alm0ra · Jul 6, 2024 · d48bea6 · d48bea6
2 parents fcbd4fd + 14983c7
commit d48bea6
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 14 deletions.
diff --git a/tests/test_aiokafka/test_async_decorators.py b/tests/test_aiokafka/test_async_decorators.py
@@ -64,14 +64,22 @@ async def test_produce_decorator(self):
     async def test_produce_twice(self):
         # subscribe to topic and get message
         self.consumer.subscribe(topics=["test"])
-        message = await self.consumer.getone()
-
-        self.assertEqual(message.value(payload=None), "test_value1")
-        self.assertEqual(message.key(), "test_key1")
 
-        message = await self.consumer.getone()
-        self.assertEqual(message.value(payload=None), "test_value")
-        self.assertEqual(message.key(), "test_key")
+        # Order unknown as partition order is not predictable
+        messages = [
+            (x.key(), x.value(payload=None))
+            for x in (
+                await self.consumer.getone(),
+                await self.consumer.getone(),
+            )
+        ]
+        self.assertCountEqual(
+            [
+                ("test_key", "test_value"),
+                ("test_key1", "test_value1"),
+            ],
+            messages,
+        )
 
         # commit message and check
         await self.consumer.commit()

diff --git a/tests/test_docrators.py b/tests/test_docrators.py
@@ -50,14 +50,22 @@ def test_produce_decorator(self):
     def test_produce_twice(self):
         # subscribe to topic and get message
         self.consumer.subscribe(topics=["test"])
-        message = self.consumer.poll()
-
-        self.assertEqual(message.value(payload=None), "test_value1")
-        self.assertEqual(message.key(), "test_key1")
 
-        message = self.consumer.poll()
-        self.assertEqual(message.value(payload=None), "test_value")
-        self.assertEqual(message.key(), "test_key")
+        # Order unknown as partition order is not predictable
+        messages = [
+            (x.key(), x.value(payload=None))
+            for x in (
+                self.consumer.poll(),
+                self.consumer.poll(),
+            )
+        ]
+        self.assertCountEqual(
+            [
+                ("test_key", "test_value"),
+                ("test_key1", "test_value1"),
+            ],
+            messages,
+        )
 
         # commit message and check
         self.consumer.commit()