From 991214ef52762ba378eacfa4cb7e2ea4d267f124 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 26 Nov 2023 09:56:32 -0600
Subject: [PATCH] Fix name logic with reconnect logic when APIConnection
 address lacks local (#728)

---
 aioesphomeapi/reconnect_logic.py |  8 ++++----
 tests/test_reconnect_logic.py    | 28 +++++++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/aioesphomeapi/reconnect_logic.py b/aioesphomeapi/reconnect_logic.py
index 2832c61..a27be75 100644
--- a/aioesphomeapi/reconnect_logic.py
+++ b/aioesphomeapi/reconnect_logic.py
@@ -19,7 +19,7 @@ from .core import (
     RequiresEncryptionAPIError,
     UnhandledAPIConnectionError,
 )
-from .util import address_is_local
+from .util import address_is_local, host_is_name_part
 from .zeroconf import ZeroconfInstanceType
 
 _LOGGER = logging.getLogger(__name__)
@@ -79,7 +79,7 @@ class ReconnectLogic(zeroconf.RecordUpdateListener):
         self.name: str | None = None
         if name:
             self.name = name
-        elif address_is_local(client.address):
+        elif host_is_name_part(client.address) or address_is_local(client.address):
             self.name = client.address.partition(".")[0]
         if self.name:
             self._cli.set_cached_name_if_unset(self.name)
@@ -229,7 +229,7 @@ class ReconnectLogic(zeroconf.RecordUpdateListener):
         if not delay:
             self._call_connect_once()
             return
-        _LOGGER.debug("Scheduling new connect attempt in %f seconds", delay)
+        _LOGGER.debug("Scheduling new connect attempt in %.2f seconds", delay)
         self._cancel_connect_timer()
         self._connect_timer = self.loop.call_at(
             self.loop.time() + delay, self._call_connect_once
@@ -303,7 +303,7 @@ class ReconnectLogic(zeroconf.RecordUpdateListener):
                 _LOGGER.info(
                     "Trying to connect to %s in the background", self._cli.log_name
                 )
-            _LOGGER.debug("Retrying %s in %d seconds", self._cli.log_name, wait_time)
+            _LOGGER.debug("Retrying %s in %.2f seconds", self._cli.log_name, wait_time)
             if wait_time:
                 # If we are waiting, start listening for mDNS records
                 self._start_zc_listen()
diff --git a/tests/test_reconnect_logic.py b/tests/test_reconnect_logic.py
index 876d391..dba1eae 100644
--- a/tests/test_reconnect_logic.py
+++ b/tests/test_reconnect_logic.py
@@ -86,13 +86,14 @@ async def test_reconnect_logic_name_from_host_and_set():
     async def on_connect() -> None:
         pass
 
-    ReconnectLogic(
+    rl = ReconnectLogic(
         client=cli,
         on_disconnect=on_disconnect,
         on_connect=on_connect,
         zeroconf_instance=get_mock_zeroconf(),
         name="mydevice",
     )
+    assert rl.name == "mydevice"
     assert cli.log_name == "mydevice.local"
 
 
@@ -145,6 +146,31 @@ async def test_reconnect_logic_name_from_name():
     assert cli.log_name == "mydevice @ 1.2.3.4"
 
 
+@pytest.mark.asyncio
+async def test_reconnect_logic_name_from_cli_address():
+    """Test that the name is set correctly from the address."""
+    cli = APIClient(
+        address="mydevice",
+        port=6052,
+        password=None,
+    )
+
+    async def on_disconnect(expected_disconnect: bool) -> None:
+        pass
+
+    async def on_connect() -> None:
+        pass
+
+    rl = ReconnectLogic(
+        client=cli,
+        on_disconnect=on_disconnect,
+        on_connect=on_connect,
+        zeroconf_instance=get_mock_zeroconf(),
+    )
+    assert cli.log_name == "mydevice"
+    assert rl.name == "mydevice"
+
+
 @pytest.mark.asyncio
 async def test_reconnect_logic_state(patchable_api_client: APIClient):
     """Test that reconnect logic state changes."""