From a0cf5d0e909e29b0b5a84683895331e05c830416 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 23 Nov 2023 18:37:25 +0100
Subject: [PATCH] Add coverage for subscribe_home_assistant_states (#677)

---
 aioesphomeapi/client.py           |  8 ++------
 aioesphomeapi/client_callbacks.py |  8 ++++++++
 tests/test_client.py              | 27 +++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/aioesphomeapi/client.py b/aioesphomeapi/client.py
index bb558e2..697945f 100644
--- a/aioesphomeapi/client.py
+++ b/aioesphomeapi/client.py
@@ -78,6 +78,7 @@ from .client_callbacks import (
     on_bluetooth_le_advertising_response,
     on_home_assistant_service_response,
     on_state_msg,
+    on_subscribe_home_assistant_state_response,
 )
 from .connection import APIConnection, ConnectionParams
 from .core import (
@@ -897,14 +898,9 @@ class APIClient:
     async def subscribe_home_assistant_states(
         self, on_state_sub: Callable[[str, str | None], None]
     ) -> None:
-        def _on_subscribe_home_assistant_state_response(
-            msg: SubscribeHomeAssistantStateResponse,
-        ) -> None:
-            on_state_sub(msg.entity_id, msg.attribute)
-
         self._get_connection().send_message_callback_response(
             SubscribeHomeAssistantStatesRequest(),
-            _on_subscribe_home_assistant_state_response,
+            partial(on_subscribe_home_assistant_state_response, on_state_sub),
             (SubscribeHomeAssistantStateResponse,),
         )
 
diff --git a/aioesphomeapi/client_callbacks.py b/aioesphomeapi/client_callbacks.py
index 5785207..483d8f4 100644
--- a/aioesphomeapi/client_callbacks.py
+++ b/aioesphomeapi/client_callbacks.py
@@ -12,6 +12,7 @@ from .api_pb2 import (  # type: ignore
     BluetoothLERawAdvertisementsResponse,
     CameraImageResponse,
     HomeassistantServiceResponse,
+    SubscribeHomeAssistantStateResponse,
 )
 from .model import (
     BluetoothLEAdvertisement,
@@ -85,3 +86,10 @@ def on_bluetooth_gatt_notify_data_response(
     """Handle a BluetoothGATTNotifyDataResponse message."""
     if address == msg.address and handle == msg.handle:
         on_bluetooth_gatt_notify(handle, bytearray(msg.data))
+
+
+def on_subscribe_home_assistant_state_response(
+    on_state_sub: Callable[[str, str | None], None],
+    msg: SubscribeHomeAssistantStateResponse,
+) -> None:
+    on_state_sub(msg.entity_id, msg.attribute)
diff --git a/tests/test_client.py b/tests/test_client.py
index 92a9657..8015585 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -48,6 +48,7 @@ from aioesphomeapi.api_pb2 import (
     NumberCommandRequest,
     SelectCommandRequest,
     SirenCommandRequest,
+    SubscribeHomeAssistantStateResponse,
     SubscribeLogsResponse,
     SwitchCommandRequest,
     TextCommandRequest,
@@ -1323,6 +1324,32 @@ async def test_subscribe_bluetooth_connections_free(
     unsub()
 
 
+@pytest.mark.asyncio
+async def test_subscribe_home_assistant_states(
+    api_client: tuple[
+        APIClient, APIConnection, asyncio.Transport, APIPlaintextFrameHelper
+    ],
+) -> None:
+    """Test subscribe_home_assistant_states."""
+    client, connection, transport, protocol = api_client
+    states = []
+
+    def on_subscribe_home_assistant_states(
+        entity_id: str, attribute: str | None
+    ) -> None:
+        states.append((entity_id, attribute))
+
+    await client.subscribe_home_assistant_states(on_subscribe_home_assistant_states)
+    await asyncio.sleep(0)
+
+    response: message.Message = SubscribeHomeAssistantStateResponse(
+        entity_id="sensor.red", attribute="any"
+    )
+    protocol.data_received(generate_plaintext_packet(response))
+
+    assert states == [("sensor.red", "any")]
+
+
 @pytest.mark.asyncio
 async def test_subscribe_logs(auth_client: APIClient) -> None:
     send = patch_response_callback(auth_client)