Speed up stopping the reconnect logic if we are at a safe place to cancel

If the connection has not established yet, we can cancel the
connection process when the reconnect logic is stopped instead
of waiting for the connection to fail. If we have already
finished connecting and started the handshake, we wait
for it to complete so we can try to disconnect cleanly

Author: bdraco

aioesphomeapi/reconnect_logic.py

@@ -317,6 +317,13 @@ async def start(self) -> None:
     async def stop(self) -> None:
         """Stop the connecting logic background task. Does not disconnect the client."""
         self._cancel_connect("Stopping")
+        if self._connect_task and self._connection_state == ReconnectLogicState.CONNECTING:
+            # If we are still establishing a connection, we can safely
+            # cancel the connect task here, otherwise we need to wait
+            # for the connect task to finish so we can gracefully
+            # disconnect.
+            self._connect_task.cancel("Stopping")
+
         async with self._connected_lock:
             self._is_stopped = True
             # Cancel again while holding the lock

tests/test_reconnect_logic.py

@@ -371,3 +371,33 @@ async def quick_connect_fail(*args, **kwargs):
         assert log_text in caplog.text
 
     await rl.stop()
+    assert rl._is_stopped is True
+    assert rl._connection_state is ReconnectLogicState.DISCONNECTED
+
+
+@pytest.mark.asyncio
+async def test_reconnect_logic_stop_callback():
+    """Test that the stop_callback stops the ReconnectLogic."""
+    cli = APIClient(
+        address="1.2.3.4",
+        port=6052,
+        password=None,
+    )
+    rl = ReconnectLogic(
+        client=cli,
+        on_disconnect=AsyncMock(),
+        on_connect=AsyncMock(),
+        zeroconf_instance=_get_mock_zeroconf(),
+        name="mydevice",
+    )
+    await rl.start()
+    assert rl._connection_state is ReconnectLogicState.DISCONNECTED
+    await asyncio.sleep(0)
+    assert rl._connection_state is ReconnectLogicState.CONNECTING
+    assert rl._is_stopped is False
+    rl.stop_callback()
+    # Wait for cancellation to propagate
+    for _ in range(4):
+        await asyncio.sleep(0)
+    assert rl._is_stopped is True
+    assert rl._connection_state is ReconnectLogicState.DISCONNECTED