From be6ead17856b7328ef1043cecd1169d73ca6746b Mon Sep 17 00:00:00 2001 From: Marcos Schroh Date: Wed, 22 Nov 2023 14:35:14 +0100 Subject: [PATCH] fix: increase total events also when using sync testing --- kstreams/test_utils/topics.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/kstreams/test_utils/topics.py b/kstreams/test_utils/topics.py index 09c6b81c..e6fe3da8 100644 --- a/kstreams/test_utils/topics.py +++ b/kstreams/test_utils/topics.py @@ -21,10 +21,7 @@ class Topic: async def put(self, event: ConsumerRecord) -> None: await self.queue.put(event) - - # keep track of the amount of events per topic partition - self.total_partition_events[event.partition] += 1 - self.total_events += 1 + self._inc_amount(event) async def get(self) -> ConsumerRecord: return await self.queue.get() @@ -33,7 +30,13 @@ def get_nowait(self) -> ConsumerRecord: return self.queue.get_nowait() def put_nowait(self, *, event: ConsumerRecord) -> None: - return self.queue.put_nowait(event) + self.queue.put_nowait(event) + self._inc_amount(event) + + def _inc_amount(self, event: ConsumerRecord) -> None: + # keep track of the amount of events per topic partition + self.total_partition_events[event.partition] += 1 + self.total_events += 1 def task_done(self) -> None: self.queue.task_done()