fix: Release lock in aenter on initial connection failure (#245); Close

#244
empicano · Aug 23, 2023 · 1f0f0b3 · 1f0f0b3
1 parent 40d4f56
commit 1f0f0b3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/aiomqtt/client.py b/aiomqtt/client.py
@@ -1039,7 +1039,11 @@ async def __aenter__(self) -> Client:
             msg = "Does not support reentrant"
             raise MqttReentrantError(msg)
         await self._lock.acquire()
-        await self.connect()
+        try:
+            await self.connect()
+        except Exception:
+            self._lock.release()
+            raise
         return self
 
     async def __aexit__(

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -13,7 +13,7 @@
 from anyio.abc import TaskStatus
 
 from aiomqtt import Client, ProtocolVersion, TLSParameters, Topic, Wildcard, Will
-from aiomqtt.error import MqttReentrantError
+from aiomqtt.error import MqttError, MqttReentrantError
 from aiomqtt.types import PayloadType
 
 pytestmark = pytest.mark.anyio
@@ -536,3 +536,11 @@ async def test_client_connecting_disconnected_done() -> None:
     client._disconnected.set_result(None)
     await client.connect()
     await client.disconnect()
+
+
+@pytest.mark.network
+async def test_client_aenter_connect_error_lock_release() -> None:
+    client = Client(hostname="aenter_connect_error_lock_release")
+    with pytest.raises(MqttError):
+        await client.__aenter__()
+    assert not client._lock.locked()