tests: fix partition tests

airtai · May 9, 2024 · aeb9115 · aeb9115
1 parent dd4e0bb
commit aeb9115
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/tests/brokers/kafka/test_test_client.py b/tests/brokers/kafka/test_test_client.py
@@ -22,13 +22,13 @@ async def test_partition_match(
         queue: str,
     ):
         broker = self.get_broker()
-      
+
         @broker.subscriber(partitions=[TopicPartition(queue, 1)])
-        async def m():
+        async def m(msg):
             pass
 
         async with self.patch_broker(broker) as br:
-            await  br.publish("hello", queue)
+            await br.publish("hello", queue)
 
             m.mock.assert_called_once_with("hello")
 
@@ -37,13 +37,13 @@ async def test_partition_match_exect(
         queue: str,
     ):
         broker = self.get_broker()
-      
+
         @broker.subscriber(partitions=[TopicPartition(queue, 1)])
-        async def m():
+        async def m(msg):
             pass
 
         async with self.patch_broker(broker) as br:
-            await test_broker.publish("hello", queue, partition=1)
+            await br.publish("hello", queue, partition=1)
 
             m.mock.assert_called_once_with("hello")
 
@@ -54,13 +54,13 @@ async def test_partition_missmatch(
         broker = self.get_broker()
 
         @broker.subscriber(partitions=[TopicPartition(queue, 1)])
-        async def m():
+        async def m(msg):
             pass
 
         @broker.subscriber(queue)
-        async def m2():
+        async def m2(msg):
             pass
-        
+
         async with self.patch_broker(broker) as br:
             await br.publish("hello", queue, partition=2)