Repository: esphome/aioesphomeapi (language: python)

Issue:
Wait for handshake to finish to allow for clean disconnect when stopping reconnect logic
There was a very rare race where we could disconnect in the middle of the handshake if we had just finished connecting to the ESP which would prevent the disconnect request from being seen by the ESP

To fix this we now only cancel without the lock if we are still establishing the connection, otherwise try to wait for the handshake to finish so we can cleanly disconnect.

This didn't cause any visible issues besides a log message on the esp


Discussion / hints:
## [Codecov](https://app.codecov.io/gh/esphome/aioesphomeapi/pull/614?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=esphome) Report
> Merging [#614](https://app.codecov.io/gh/esphome/aioesphomeapi/pull/614?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=esphome) (6eb9ff1) into [main](https://app.codecov.io/gh/esphome/aioesphomeapi/commit/2ef9ed90260d72396c6bbd868de012cb2c4e285a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=esphome) (2ef9ed9) will **increase** coverage by `0.41%`.
> Report is 1 commits behind head on main.
> The diff coverage is `50.00%`.

> :exclamation: Current head 6eb9ff1 differs from pull request most recent head 9ade838. Consider uploading reports for the commit 9ade838 to get more accurate results

```diff
@@            Coverage Diff             @@
##             main     #614      +/-   ##
==========================================
+ Coverage   77.41%   77.83%   +0.41%     
==========================================
  Files          14       14              
  Lines        2488     2490       +2     
==========================================
+ Hits         1926     1938      +12     
+ Misses        562      552      -10     
```


| [Files](https://app.codecov.io/gh/esphome/aioesphomeapi/pull/614?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=esphome) | Coverage Δ | |
|---|---|---|
| [aioesphomeapi/reconnect\_logic.py](https://app.codecov.io/gh/esphome/aioesphomeapi/pull/614?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=esphome#diff-YWlvZXNwaG9tZWFwaS9yZWNvbm5lY3RfbG9naWMucHk=) | `87.83% <50.00%> (+1.19%)` | :arrow_up: |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/esphome/aioesphomeapi/pull/614/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=esphome)

:mega: Codecov offers a browser extension for seamless coverage viewing on GitHub. Try it in [Chrome](https://chrome.google.com/webstore/detail/codecov/gedikamndpbemklijjkncpnolildpbgo) or [Firefox](https://addons.mozilla.org/en-US/firefox/addon/codecov/) today!


Fix patch:
diff --git a/aioesphomeapi/reconnect_logic.py b/aioesphomeapi/reconnect_logic.py
index 154c7d17..a665ae37 100644
--- a/aioesphomeapi/reconnect_logic.py
+++ b/aioesphomeapi/reconnect_logic.py
@@ -316,7 +316,7 @@ async def start(self) -> None:
 
     async def stop(self) -> None:
         """Stop the connecting logic background task. Does not disconnect the client."""
-        if self._connection_state == ReconnectLogicState.CONNECTING:
+        if self._connection_state in NOT_YET_CONNECTED_STATES:
             # 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


Test patch:
diff --git a/tests/test_reconnect_logic.py b/tests/test_reconnect_logic.py
index e3f7842d..3d814319 100644
--- a/tests/test_reconnect_logic.py
+++ b/tests/test_reconnect_logic.py
@@ -401,3 +401,52 @@ async def test_reconnect_logic_stop_callback():
         await asyncio.sleep(0)
     assert rl._is_stopped is True
     assert rl._connection_state is ReconnectLogicState.DISCONNECTED
+
+
+@pytest.mark.asyncio
+async def test_reconnect_logic_stop_callback_waits_for_handshake():
+    """Test that the stop_callback waits for a handshake."""
+
+    class PatchableAPIClient(APIClient):
+        pass
+
+    cli = PatchableAPIClient(
+        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",
+    )
+    assert rl._connection_state is ReconnectLogicState.DISCONNECTED
+
+    async def slow_connect_fail(*args, **kwargs):
+        await asyncio.sleep(10)
+        raise APIConnectionError
+
+    with patch.object(cli, "start_connection"), patch.object(
+        cli, "finish_connection", side_effect=slow_connect_fail
+    ):
+        await rl.start()
+        for _ in range(3):
+            await asyncio.sleep(0)
+
+    assert rl._connection_state is ReconnectLogicState.HANDSHAKING
+    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 False
+    assert rl._connection_state is ReconnectLogicState.HANDSHAKING
+
+    rl._cancel_connect("forced cancel in test")
+    # 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