From 034730222276780c0f3b955ba779a8211ce594a7 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 23 Nov 2023 17:49:24 +0100
Subject: [PATCH] Add coverage for bluetooth connection free responses (#673)

---
 aioesphomeapi/client.py | 21 ++++++++++++---------
 tests/test_client.py    | 25 +++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/aioesphomeapi/client.py b/aioesphomeapi/client.py
index 94909e0..2ada885 100644
--- a/aioesphomeapi/client.py
+++ b/aioesphomeapi/client.py
@@ -618,20 +618,23 @@ class APIClient:
 
         return unsub
 
+    def _on_bluetooth_connections_free_response(
+        self,
+        on_bluetooth_connections_free_update: Callable[[int, int], None],
+        msg: BluetoothConnectionsFreeResponse,
+    ) -> None:
+        on_bluetooth_connections_free_update(msg.free, msg.limit)
+
     async def subscribe_bluetooth_connections_free(
         self, on_bluetooth_connections_free_update: Callable[[int, int], None]
     ) -> Callable[[], None]:
-        msg_types = (BluetoothConnectionsFreeResponse,)
-
-        def _on_bluetooth_connections_free_response(
-            msg: BluetoothConnectionsFreeResponse,
-        ) -> None:
-            on_bluetooth_connections_free_update(msg.free, msg.limit)
-
         return self._get_connection().send_message_callback_response(
             SubscribeBluetoothConnectionsFreeRequest(),
-            _on_bluetooth_connections_free_response,
-            msg_types,
+            partial(
+                self._on_bluetooth_connections_free_response,
+                on_bluetooth_connections_free_update,
+            ),
+            (BluetoothConnectionsFreeResponse,),
         )
 
     def _handle_timeout(self, fut: asyncio.Future[None]) -> None:
diff --git a/tests/test_client.py b/tests/test_client.py
index 497bb0a..92a9657 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -12,6 +12,7 @@ from aioesphomeapi._frame_helper.plain_text import APIPlaintextFrameHelper
 from aioesphomeapi.api_pb2 import (
     AlarmControlPanelCommandRequest,
     BinarySensorStateResponse,
+    BluetoothConnectionsFreeResponse,
     BluetoothDeviceClearCacheResponse,
     BluetoothDeviceConnectionResponse,
     BluetoothDevicePairingResponse,
@@ -1298,6 +1299,30 @@ async def test_subscribe_bluetooth_le_raw_advertisements(
     unsub()
 
 
+@pytest.mark.asyncio
+async def test_subscribe_bluetooth_connections_free(
+    api_client: tuple[
+        APIClient, APIConnection, asyncio.Transport, APIPlaintextFrameHelper
+    ],
+) -> None:
+    """Test subscribe_bluetooth_connections_free."""
+    client, connection, transport, protocol = api_client
+    connections = []
+
+    def on_bluetooth_connections_free(free: int, limit: int) -> None:
+        connections.append((free, limit))
+
+    unsub = await client.subscribe_bluetooth_connections_free(
+        on_bluetooth_connections_free
+    )
+    await asyncio.sleep(0)
+    response: message.Message = BluetoothConnectionsFreeResponse(free=2, limit=3)
+    protocol.data_received(generate_plaintext_packet(response))
+
+    assert connections == [(2, 3)]
+    unsub()
+
+
 @pytest.mark.asyncio
 async def test_subscribe_logs(auth_client: APIClient) -> None:
     send = patch_response_callback(auth_client)