Add context manager regression test (#1130)

Signed-off-by: Callan Gray <[email protected]>

Author: calgray

tests/test_consumer.py

@@ -2,6 +2,7 @@
 import gc
 import json
 import time
+import warnings
 from contextlib import contextmanager
 from unittest import mock
 
@@ -188,6 +189,36 @@ async def test_consumer_context_manager(self):
                 raise ValueError
         assert consumer._closed
 
+    @run_until_complete
+    async def test_consumer_context_manager_start_error(self):
+        for target, group_id in [
+            ("aiokafka.consumer.consumer.AIOKafkaClient.bootstrap", None),
+            ("aiokafka.consumer.consumer.Fetcher.__init__", None),
+            ("aiokafka.consumer.consumer.NoGroupCoordinator.__init__", None),
+            (
+                "aiokafka.consumer.consumer.GroupCoordinator.__init__",
+                f"group-{self.id()}",
+            ),
+        ]:
+            consumer = AIOKafkaConsumer(
+                self.topic,
+                group_id=group_id,
+                bootstrap_servers=self.hosts,
+                enable_auto_commit=False,
+                auto_offset_reset="earliest",
+            )
+
+            # make consumer.start() raise
+            with mock.patch(target, side_effect=RuntimeError("error")):
+                with self.assertRaises(RuntimeError):
+                    async with consumer:
+                        self.fail(f"{target} did not raise")
+
+            # should not require calling consumer.close()
+            with warnings.catch_warnings(record=True) as w:
+                del consumer
+                self.assertEqual(len(w), 0, f"{target} got unexpected warnings: {w}")
+
     @run_until_complete
     async def test_consumer_api_version(self):
         await self.send_messages(0, list(range(10)))

tests/test_producer.py

@@ -2,6 +2,7 @@
 import gc
 import json
 import time
+import warnings
 import weakref
 from unittest import mock
 
@@ -179,6 +180,25 @@ async def test_producer_context_manager(self):
                 raise ValueError()
         assert producer._closed
 
+    @run_until_complete
+    async def test_producer_context_manager_start_error(self):
+        for target in [
+            "aiokafka.producer.producer.AIOKafkaClient.bootstrap",
+            "aiokafka.producer.producer.Sender.start",
+        ]:
+            producer = AIOKafkaProducer(bootstrap_servers=self.hosts)
+
+            # make producer.start() raise
+            with mock.patch(target, side_effect=RuntimeError("error")):
+                with self.assertRaises(RuntimeError):
+                    async with producer:
+                        self.fail(f"{target} did not raise")
+
+            # should not require calling producer.close()
+            with warnings.catch_warnings(record=True) as w:
+                del producer
+                self.assertEqual(len(w), 0, f"Unexpected warnings: {w}")
+
     @run_until_complete
     async def test_producer_send_noack(self):
         producer = AIOKafkaProducer(bootstrap_servers=self.hosts, acks=0)