test: Add tests for _disconnect future handling in __aexit__

aiomqtt/client.py

@@ -491,18 +491,16 @@ async def connect(self, *, timeout: float | None = None) -> None:
         self._logger.warning(
             "The manual `connect` and `disconnect` methods are deprecated and will be"
             " removed in a future version. The preferred way to connect and disconnect"
-            " the client is to use the context manager interface via `async with`."
-            " In case you're forced to connect and disconnect manually, you can call"
-            " the context manager's `__aenter__` and `__aexit__` methods directly as a"
-            " workaround instead. `__aenter__` is equivalent to `connect`. `__aexit__`"
-            " is equivalent to `disconnect` except that it forces disconnection instead"
-            " of throwing an exception in case the client cannot disconnect normally."
+            " the client is to use the context manager interface via `async with`. In"
+            " case your use case needs to connect and disconnect manually, you can call"
+            " the context manager's `__aenter__` and `__aexit__` methods as an escape"
+            " hatch instead. `__aenter__` is equivalent to `connect`. `__aexit__` is"
+            " equivalent to `disconnect` except that it forces disconnection instead"
+            " of throwing an exception in case the client cannot disconnect cleanly."
             " `__aexit__` expects three arguments: `exc_type`, `exc`, and `tb`. These"
             " arguments describe the exception that caused the context manager to exit,"
             " if any. You can pass `None` to all of these arguments in a manual call to"
-            " `__aexit__`. Note that with calling `__aenter__` and `__aexit__`"
-            " directly, you are using internal functionality that may break without"
-            " prior notice."
+            " `__aexit__`."
         )
         try:
             loop = asyncio.get_running_loop()
@@ -548,18 +546,16 @@ async def disconnect(self, *, timeout: float | None = None) -> None:
         self._logger.warning(
             "The manual `connect` and `disconnect` methods are deprecated and will be"
             " removed in a future version. The preferred way to connect and disconnect"
-            " the client is to use the context manager interface via `async with`."
-            " In case you're forced to connect and disconnect manually, you can call"
-            " the context manager's `__aenter__` and `__aexit__` methods directly as a"
-            " workaround instead. `__aenter__` is equivalent to `connect`. `__aexit__`"
-            " is equivalent to `disconnect` except that it forces disconnection instead"
-            " of throwing an exception in case the client cannot disconnect normally."
+            " the client is to use the context manager interface via `async with`. In"
+            " case your use case needs to connect and disconnect manually, you can call"
+            " the context manager's `__aenter__` and `__aexit__` methods as an escape"
+            " hatch instead. `__aenter__` is equivalent to `connect`. `__aexit__` is"
+            " equivalent to `disconnect` except that it forces disconnection instead"
+            " of throwing an exception in case the client cannot disconnect cleanly."
             " `__aexit__` expects three arguments: `exc_type`, `exc`, and `tb`. These"
             " arguments describe the exception that caused the context manager to exit,"
             " if any. You can pass `None` to all of these arguments in a manual call to"
-            " `__aexit__`. Note that with calling `__aenter__` and `__aexit__`"
-            " directly, you are using internal functionality that may break without"
-            " prior notice."
+            " `__aexit__`."
         )
         if self._early_out_on_disconnected():
             return
@@ -1103,7 +1099,12 @@ async def __aexit__(
         tb: TracebackType | None,
     ) -> None:
         """Disconnect from the broker."""
-        if self._early_out_on_disconnected():
+        if self._disconnected.done():
+            # Return early if the client is already disconnected
+            self._lock.release()
+            if (exc := self._disconnected.exception()) is not None:
+                # If the disconnect wasn't intentional, raise the error that caused it
+                raise exc
             return
         # Try to gracefully disconnect from the broker
         rc = self._client.disconnect()

scripts/test

@@ -6,4 +6,4 @@ set -o errexit -o pipefail -o nounset
 cd "$(dirname "$0")/.."
 
 # Run tests with pytest
-poetry run pytest --failed-first --strict-config --strict-markers --verbosity=2 --cov=aiomqtt --cov-report=xml --junitxml=reports/pytest.xml tests
+poetry run pytest --failed-first --verbosity=2 --cov=aiomqtt --cov-report=term-missing --cov-report=xml --junitxml=reports/pytest.xml --strict-config --strict-markers tests

tests/test_client.py

@@ -88,15 +88,6 @@ async def test_topic_matches() -> None:
     assert not topic.matches("$test/group/a/b/c")
 
 
-async def test__aexit__() -> None:
-    """Test that it's possible to call __aexit__ without prior __aenter__.
-
-    This should also cover the case of an unsucessful __aenter__.
-    """
-    client = Client(HOSTNAME)
-    await client.__aexit__(None, None, None)
-
-
 @pytest.mark.network
 async def test_multiple_messages_generators() -> None:
     """Test that multiple Client.messages() generators can be used at the same time."""
@@ -415,21 +406,20 @@ async def test_client_no_pending_calls_warnings_with_max_concurrent_outgoing_cal
 
 @pytest.mark.network
 async def test_client_not_reentrant() -> None:
+    """Test that the client raises an error when we try to reenter."""
     client = Client(HOSTNAME)
-
     with pytest.raises(MqttReentrantError):  # noqa: PT012
         async with client:
             async with client:
-                ...
+                pass
 
 
 @pytest.mark.network
 async def test_client_reusable() -> None:
+    """Test that an instance of the client context manager can be reused."""
     client = Client(HOSTNAME)
-
     async with client:
         await client.publish("task/a", "task_a")
-
     async with client:
         await client.publish("task/b", "task_b")
 
@@ -549,8 +539,35 @@ async def test_client_connecting_disconnected_done() -> None:
 
 
 @pytest.mark.network
-async def test_client_aenter_connect_error_lock_release() -> None:
+async def test_client_aenter_error_lock_release() -> None:
+    """Test that the client's reusability lock is released on error in __aenter__."""
     client = Client(hostname="aenter_connect_error_lock_release")
     with pytest.raises(MqttError):
         await client.__aenter__()
     assert not client._lock.locked()
+
+
+@pytest.mark.network
+async def test_aexit_without_prior_aenter() -> None:
+    """Test that __aexit__ without prior (or unsuccessful) __aenter__ runs cleanly."""
+    client = Client(HOSTNAME)
+    await client.__aexit__(None, None, None)
+
+
+@pytest.mark.network
+async def test_aexit_client_is_already_disconnected_sucess() -> None:
+    """Test that __aexit__ exits cleanly if client is already cleanly disconnected."""
+    client = Client(HOSTNAME)
+    await client.__aenter__()
+    client._disconnected.set_result(None)
+    await client.__aexit__(None, None, None)
+
+
+@pytest.mark.network
+async def test_aexit_client_is_already_disconnected_failure() -> None:
+    """Test that __aexit__ reraises if client is already disconnected with an error."""
+    client = Client(HOSTNAME)
+    await client.__aenter__()
+    client._disconnected.set_exception(RuntimeError)
+    with pytest.raises(RuntimeError):
+        await client.__aexit__(None, None, None)