Fix flaky test_max_poll_interval_ms

aio-libs · Jul 1, 2023 · bdf014b · bdf014b
1 parent 5ce5649
commit bdf014b
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/tests/test_consumer.py b/tests/test_consumer.py
@@ -2018,9 +2018,15 @@ async def on_partitions_assigned(self, assigned):
         consumer1.subscribe([self.topic], listener=listener1)
         consumer2.subscribe([self.topic], listener=listener2)
 
-        # Make sure we rebalanced and ready for processing each of it's part
-        await listener1.assignment_ready.wait()
-        await listener2.assignment_ready.wait()
+        for _ in range(5):
+            # Make sure we rebalanced and ready for processing each of it's part
+            await listener1.assignment_ready.wait()
+            await listener2.assignment_ready.wait()
+
+            # Check the first is still ready to avoid flakiness
+            if listener1.assignment_ready.is_set():
+                break
+
         self.assertTrue(consumer1.assignment())
         self.assertTrue(consumer2.assignment())