diff --git a/tests/integration/fixtures/processing.py b/tests/integration/fixtures/processing.py index 1bdef35e23..22b71330f0 100644 --- a/tests/integration/fixtures/processing.py +++ b/tests/integration/fixtures/processing.py @@ -156,7 +156,7 @@ def __init__(self, consumer, options, topic_name, timeout=None): self.consumer = consumer self.test_producer = kafka_producer(options) self.topic_name = topic_name - self.timeout = timeout or 2 + self.timeout = timeout or 5 # Connect to the topic and poll a first test message. # First poll takes forever, the next ones are fast. @@ -246,13 +246,15 @@ def get_outcome(self, timeout=None): assert len(outcomes) == 1, "More than one outcome was consumed" return outcomes[0] - def assert_rate_limited(self, reason, key_id=None, categories=None, quantity=None): + def assert_rate_limited( + self, reason, key_id=None, categories=None, quantity=None, timeout=1 + ): if categories is None: - outcome = self.get_outcome() + outcome = self.get_outcome(timeout=timeout) assert isinstance(outcome["category"], int) outcomes = [outcome] else: - outcomes = self.get_outcomes() + outcomes = self.get_outcomes(timeout=timeout) expected = {category_value(category) for category in categories} actual = {outcome["category"] for outcome in outcomes} assert actual == expected, (actual, expected)