From f7697de129c28052f2ebd2ae005e3c69ae66943d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20B=C3=B6hm?= <felix@felixboehm.dev>
Date: Sun, 17 Mar 2024 14:33:42 -0300
Subject: [PATCH] test: Add test for state reset on negative CONNACK

---
 aiomqtt/client.py    |  5 ++++-
 tests/test_client.py | 18 ++++++++++++++----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/aiomqtt/client.py b/aiomqtt/client.py
index 79c559f..8d64de4 100644
--- a/aiomqtt/client.py
+++ b/aiomqtt/client.py
@@ -531,6 +531,7 @@ def _on_connect(  # noqa: PLR0913
         if rc == mqtt.CONNACK_ACCEPTED:
             self._connected.set_result(None)
         else:
+            # We received a negative CONNACK response
             self._connected.set_exception(MqttConnectError(rc))
 
     def _on_disconnect(
@@ -572,6 +573,7 @@ def _on_subscribe(  # noqa: PLR0913
         granted_qos: tuple[int] | list[mqtt.ReasonCodes],
         properties: mqtt.Properties | None = None,
     ) -> None:
+        """Called when we receive a SUBACK message from the broker."""
         try:
             fut = self._pending_subscribes.pop(mid)
             if not fut.done():
@@ -589,6 +591,7 @@ def _on_unsubscribe(  # noqa: PLR0913
         properties: mqtt.Properties | None = None,
         reason_codes: list[mqtt.ReasonCodes] | mqtt.ReasonCodes | None = None,
     ) -> None:
+        """Called when we receive an UNSUBACK message from the broker."""
         try:
             self._pending_unsubscribes.pop(mid).set()
         except KeyError:
@@ -704,7 +707,7 @@ async def __aenter__(self) -> Self:
         try:
             await self._wait_for(self._connected, timeout=None)
         except MqttError:
-            # Reset internal state if the connection attempt times out
+            # Reset state if connection attempt times out or CONNACK returns negative
             self._lock.release()
             self._connected = asyncio.Future()
             raise
diff --git a/tests/test_client.py b/tests/test_client.py
index 4047904..333a882 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -328,8 +328,8 @@ async def task_a_publisher() -> None:
 
 
 @pytest.mark.network
-async def test_aenter_lock_release_connection_failure() -> None:
-    """Test that reusability lock is released on connection failure in ``aenter``."""
+async def test_aenter_state_reset_connect_failure() -> None:
+    """Test that internal state is reset on CONNECT failure in ``aenter``."""
     client = Client(hostname="invalid")
     with pytest.raises(MqttError):
         await client.__aenter__()
@@ -338,8 +338,8 @@ async def test_aenter_lock_release_connection_failure() -> None:
 
 
 @pytest.mark.network
-async def test_aenter_lock_release_connection_timeout() -> None:
-    """Test that reusability lock is released on connection timeout in ``aenter``."""
+async def test_aenter_state_reset_connack_timeout() -> None:
+    """Test that internal state is reset on CONNACK timeout in ``aenter``."""
     client = Client(HOSTNAME, timeout=0)
     with pytest.raises(MqttError):
         await client.__aenter__()
@@ -347,6 +347,16 @@ async def test_aenter_lock_release_connection_timeout() -> None:
     assert not client._connected.done()
 
 
+@pytest.mark.network
+async def test_aenter_state_reset_connack_negative() -> None:
+    """Test that internal state is reset on negative CONNACK in ``aenter``."""
+    client = Client(HOSTNAME, username="invalid")
+    with pytest.raises(MqttError):
+        await client.__aenter__()
+    assert not client._lock.locked()
+    assert not client._connected.done()
+
+
 @pytest.mark.network
 async def test_aexit_without_prior_aenter() -> None:
     """Test that ``aexit`` without prior (or unsuccessful) ``aenter`` runs cleanly."""