Repository: empicano/aiomqtt (language: python)

Issue:
Exceptions on __aenter__ not handled.
We are using aiomqtt to connect to AWS IoT Core with client certificate and just in time registration, this means first connection to broker will fail with an `Operation Timed out` exception from line `client.py: __aenter__().await self._wait_for(self._connected, timeout=None)`. We have the following retry logic, but we continuously get  "MqttReentrantError" exception for following iterations since `self._lock` is not released and `self._connected` future is not re-initialized for the second call.

```python
while True:
    try:
       async with client:
          ...
```

The following change in `client.py: __aenter__()` solved the problem.
```diff
-        await self._wait_for(self._connected, timeout=None)
+        try:
+            await self._wait_for(self._connected, timeout=None)
+        except Exception as e:
+            self._lock.release()
+            self._connected = asyncio.Future()
+            raise e
```


Discussion / hints:
Following since I was debugging this with John.
Feel free to reach out to either of us if we can help with a reproduction!

Fix patch:
diff --git a/aiomqtt/client.py b/aiomqtt/client.py
index 35f8536..79c559f 100644
--- a/aiomqtt/client.py
+++ b/aiomqtt/client.py
@@ -701,7 +701,13 @@ async def __aenter__(self) -> Self:
         except (OSError, mqtt.WebsocketConnectionError) as exc:
             self._lock.release()
             raise MqttError(str(exc)) from None
-        await self._wait_for(self._connected, timeout=None)
+        try:
+            await self._wait_for(self._connected, timeout=None)
+        except MqttError:
+            # Reset internal state if the connection attempt times out
+            self._lock.release()
+            self._connected = asyncio.Future()
+            raise
         # Reset `_disconnected` if it's already in completed state after connecting
         if self._disconnected.done():
             self._disconnected = asyncio.Future()


Test patch:
diff --git a/tests/test_client.py b/tests/test_client.py
index 9136988..4047904 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -328,12 +328,23 @@ async def task_a_publisher() -> None:
 
 
 @pytest.mark.network
-async def test_aenter_error_lock_release() -> None:
-    """Test that the client's reusability lock is released on error in ``aenter``."""
+async def test_aenter_lock_release_connection_failure() -> None:
+    """Test that reusability lock is released on connection failure in ``aenter``."""
     client = Client(hostname="invalid")
     with pytest.raises(MqttError):
         await client.__aenter__()
     assert not client._lock.locked()
+    assert not client._connected.done()
+
+
+@pytest.mark.network
+async def test_aenter_lock_release_connection_timeout() -> None:
+    """Test that reusability lock is released on connection timeout in ``aenter``."""
+    client = Client(HOSTNAME, timeout=0)
+    with pytest.raises(MqttError):
+        await client.__aenter__()
+    assert not client._lock.locked()
+    assert not client._connected.done()
 
 
 @pytest.mark.network