Skip to content

Commit

Permalink
test: Add test for state reset on negative CONNACK
Browse files Browse the repository at this point in the history
  • Loading branch information
empicano committed Mar 17, 2024
1 parent 063244f commit f7697de
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 5 deletions.
5 changes: 4 additions & 1 deletion aiomqtt/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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():
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand Down
18 changes: 14 additions & 4 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__()
Expand All @@ -338,15 +338,25 @@ 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__()
assert not client._lock.locked()
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."""
Expand Down

0 comments on commit f7697de

Please sign in to comment.