From 4fe46f5aa0174fa91289601b3efe3eb88829a5ef Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 25 Jun 2023 21:45:36 -0500
Subject: [PATCH] Propagate if disconnect was expected to disconnect callback
 (#448)

---
 aioesphomeapi/log_reader.py      | 4 +++-
 aioesphomeapi/reconnect_logic.py | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/aioesphomeapi/log_reader.py b/aioesphomeapi/log_reader.py
index 8519a9a..dd6b2a1 100644
--- a/aioesphomeapi/log_reader.py
+++ b/aioesphomeapi/log_reader.py
@@ -59,7 +59,9 @@ async def main(argv: List[str]) -> None:
         except APIConnectionError:
             await cli.disconnect()
 
-    async def on_disconnect() -> None:
+    async def on_disconnect(  # pylint: disable=unused-argument
+        expected_disconnect: bool,
+    ) -> None:
         _LOGGER.warning("Disconnected from API")
 
     logic = ReconnectLogic(
diff --git a/aioesphomeapi/reconnect_logic.py b/aioesphomeapi/reconnect_logic.py
index 1129158..1f79f60 100644
--- a/aioesphomeapi/reconnect_logic.py
+++ b/aioesphomeapi/reconnect_logic.py
@@ -33,7 +33,7 @@ class ReconnectLogic(zeroconf.RecordUpdateListener):
         *,
         client: APIClient,
         on_connect: Callable[[], Awaitable[None]],
-        on_disconnect: Callable[[], Awaitable[None]],
+        on_disconnect: Callable[[bool], Awaitable[None]],
         zeroconf_instance: "zeroconf.Zeroconf",
         name: Optional[str] = None,
         on_connect_error: Optional[Callable[[Exception], Awaitable[None]]] = None,
@@ -86,7 +86,7 @@ class ReconnectLogic(zeroconf.RecordUpdateListener):
         )
 
         # Run disconnect hook
-        await self._on_disconnect_cb()
+        await self._on_disconnect_cb(expected_disconnect)
 
         async with self._connected_lock:
             self._connected = False