Wait for handshake to finish to allow for clean disconnect when stopping reconnect logic (#614)

aioesphomeapi/reconnect_logic.py

@@ -316,7 +316,13 @@ class ReconnectLogic(zeroconf.RecordUpdateListener):
 
     async def stop(self) -> None:
         """Stop the connecting logic background task. Does not disconnect the client."""
-        self._cancel_connect("Stopping")
+        if 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._cancel_connect("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 test_reconnect_zeroconf(
         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