From 33966938f27cb45a3966cb3777ca4df78295f819 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 23 Nov 2023 17:20:56 +0100
Subject: [PATCH] Add coverage for bluetooth advertising (#670)

---
 aioesphomeapi/client.py | 33 +++++++++-------
 tests/test_client.py    | 86 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+), 14 deletions(-)

diff --git a/aioesphomeapi/client.py b/aioesphomeapi/client.py
index 4808536..c4a9a01 100644
--- a/aioesphomeapi/client.py
+++ b/aioesphomeapi/client.py
@@ -563,19 +563,23 @@ class APIClient:
 
         return resp[0]
 
+    def _on_bluetooth_le_advertising_response(
+        self,
+        on_bluetooth_le_advertisement: Callable[[BluetoothLEAdvertisement], None],
+        msg: BluetoothLEAdvertisementResponse,
+    ) -> None:
+        on_bluetooth_le_advertisement(BluetoothLEAdvertisement.from_pb(msg))  # type: ignore[misc]
+
     async def subscribe_bluetooth_le_advertisements(
         self, on_bluetooth_le_advertisement: Callable[[BluetoothLEAdvertisement], None]
     ) -> Callable[[], None]:
         msg_types = (BluetoothLEAdvertisementResponse,)
-
-        def _on_bluetooth_le_advertising_response(
-            msg: BluetoothLEAdvertisementResponse,
-        ) -> None:
-            on_bluetooth_le_advertisement(BluetoothLEAdvertisement.from_pb(msg))  # type: ignore[misc]
-
         unsub_callback = self._get_connection().send_message_callback_response(
             SubscribeBluetoothLEAdvertisementsRequest(flags=0),
-            _on_bluetooth_le_advertising_response,
+            partial(
+                self._on_bluetooth_le_advertising_response,
+                on_bluetooth_le_advertisement,
+            ),
             msg_types,
         )
 
@@ -588,21 +592,22 @@ class APIClient:
 
         return unsub
 
+    def _on_ble_raw_advertisement_response(
+        self,
+        on_advertisements: Callable[[list[BluetoothLERawAdvertisement]], None],
+        msg: BluetoothLERawAdvertisementsResponse,
+    ) -> None:
+        on_advertisements(msg.advertisements)
+
     async def subscribe_bluetooth_le_raw_advertisements(
         self, on_advertisements: Callable[[list[BluetoothLERawAdvertisement]], None]
     ) -> Callable[[], None]:
         msg_types = (BluetoothLERawAdvertisementsResponse,)
-
-        def _on_ble_raw_advertisement_response(
-            data: BluetoothLERawAdvertisementsResponse,
-        ) -> None:
-            on_advertisements(data.advertisements)
-
         unsub_callback = self._get_connection().send_message_callback_response(
             SubscribeBluetoothLEAdvertisementsRequest(
                 flags=BluetoothProxySubscriptionFlag.RAW_ADVERTISEMENTS
             ),
-            _on_ble_raw_advertisement_response,
+            partial(self._on_ble_raw_advertisement_response, on_advertisements),
             msg_types,
         )
 
diff --git a/tests/test_client.py b/tests/test_client.py
index faed270..497bb0a 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -24,6 +24,9 @@ from aioesphomeapi.api_pb2 import (
     BluetoothGATTReadResponse,
     BluetoothGATTService,
     BluetoothGATTWriteResponse,
+    BluetoothLEAdvertisementResponse,
+    BluetoothLERawAdvertisement,
+    BluetoothLERawAdvertisementsResponse,
     ButtonCommandRequest,
     CameraImageRequest,
     CameraImageResponse,
@@ -63,6 +66,7 @@ from aioesphomeapi.model import (
 )
 from aioesphomeapi.model import BluetoothGATTService as BluetoothGATTServiceModel
 from aioesphomeapi.model import (
+    BluetoothLEAdvertisement,
     CameraState,
     ClimateFanMode,
     ClimateMode,
@@ -1212,6 +1216,88 @@ async def test_bluetooth_gatt_start_notify_fails(
     )
 
 
+@pytest.mark.asyncio
+async def test_subscribe_bluetooth_le_advertisements(
+    api_client: tuple[
+        APIClient, APIConnection, asyncio.Transport, APIPlaintextFrameHelper
+    ],
+) -> None:
+    """Test subscribe_bluetooth_le_advertisements."""
+    client, connection, transport, protocol = api_client
+    advs = []
+
+    def on_bluetooth_le_advertisements(adv: BluetoothLEAdvertisement) -> None:
+        advs.append(adv)
+
+    unsub = await client.subscribe_bluetooth_le_advertisements(
+        on_bluetooth_le_advertisements
+    )
+    await asyncio.sleep(0)
+    response: message.Message = BluetoothLEAdvertisementResponse(
+        address=1234,
+        name=b"mydevice",
+        rssi=-50,
+        service_uuids=["1234"],
+        service_data={},
+        manufacturer_data={},
+        address_type=1,
+    )
+    protocol.data_received(generate_plaintext_packet(response))
+
+    assert advs == [
+        BluetoothLEAdvertisement(
+            address=1234,
+            name="mydevice",
+            rssi=-50,
+            service_uuids=["000034-0000-1000-8000-00805f9b34fb"],
+            manufacturer_data={},
+            service_data={},
+            address_type=1,
+        )
+    ]
+    unsub()
+
+
+@pytest.mark.asyncio
+async def test_subscribe_bluetooth_le_raw_advertisements(
+    api_client: tuple[
+        APIClient, APIConnection, asyncio.Transport, APIPlaintextFrameHelper
+    ],
+) -> None:
+    """Test subscribe_bluetooth_le_raw_advertisements."""
+    client, connection, transport, protocol = api_client
+    adv_groups = []
+
+    def on_raw_bluetooth_le_advertisements(
+        advs: list[BluetoothLERawAdvertisementsResponse],
+    ) -> None:
+        adv_groups.append(advs)
+
+    unsub = await client.subscribe_bluetooth_le_raw_advertisements(
+        on_raw_bluetooth_le_advertisements
+    )
+    await asyncio.sleep(0)
+
+    response: message.Message = BluetoothLERawAdvertisementsResponse(
+        advertisements=[
+            BluetoothLERawAdvertisement(
+                address=1234,
+                rssi=-50,
+                address_type=1,
+                data=b"1234",
+            )
+        ]
+    )
+    protocol.data_received(generate_plaintext_packet(response))
+    assert len(adv_groups) == 1
+    first_adv = adv_groups[0][0]
+    assert first_adv.address == 1234
+    assert first_adv.rssi == -50
+    assert first_adv.address_type == 1
+    assert first_adv.data == b"1234"
+    unsub()
+
+
 @pytest.mark.asyncio
 async def test_subscribe_logs(auth_client: APIClient) -> None:
     send = patch_response_callback(auth_client)