Fix AIOKafkaProducer context manager (#729)

aio-libs · Mar 17, 2021 · 4c9f827 · 4c9f827
1 parent 5d9af4d
commit 4c9f827
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/aiokafka/producer/producer.py b/aiokafka/producer/producer.py
@@ -561,6 +561,7 @@ async def send_offsets_to_transaction(self, offsets, group_id):
 
     async def __aenter__(self):
         await self.start()
+        return self
 
     async def __aexit__(self, exc_type, exc, tb):
         await self.stop()

diff --git a/tests/test_producer.py b/tests/test_producer.py
@@ -153,7 +153,8 @@ def test_create_producer_no_running_loop(self):
     async def test_producer_context_manager(self):
         producer = AIOKafkaProducer(
             bootstrap_servers=self.hosts)
-        async with producer:
+        async with producer as prod:
+            assert prod is producer
             assert producer._sender._sender_task is not None
             await producer.send(self.topic, b'value', key=b'KEY')
         assert producer._closed
@@ -162,7 +163,8 @@ async def test_producer_context_manager(self):
         producer = AIOKafkaProducer(
             bootstrap_servers=self.hosts)
         with pytest.raises(ValueError):
-            async with producer:
+            async with producer as prod:
+                assert prod is producer
                 assert producer._sender._sender_task is not None
                 await producer.send(self.topic, b'value', key=b'KEY')
                 raise ValueError()