From cf2fd3c92aa779da89b01eb70c05dbf24b3afcfd Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 23 Nov 2023 16:17:41 +0100
Subject: [PATCH] Add test coverage for bluetooth GATT notify (#668)

---
 tests/test_client.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/tests/test_client.py b/tests/test_client.py
index d4d1472..a3abe3f 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -18,6 +18,8 @@ from aioesphomeapi.api_pb2 import (
     BluetoothGATTErrorResponse,
     BluetoothGATTGetServicesDoneResponse,
     BluetoothGATTGetServicesResponse,
+    BluetoothGATTNotifyDataResponse,
+    BluetoothGATTNotifyResponse,
     BluetoothGATTReadResponse,
     BluetoothGATTService,
     BluetoothGATTWriteResponse,
@@ -1128,6 +1130,45 @@ async def test_bluetooth_gatt_get_services_errors(
         await services_task
 
 
+@pytest.mark.xfail(reason="There is a race condition here")
+@pytest.mark.asyncio
+async def test_bluetooth_gatt_start_notify(
+    api_client: tuple[
+        APIClient, APIConnection, asyncio.Transport, APIPlaintextFrameHelper
+    ],
+) -> None:
+    """Test bluetooth_gatt_start_notify."""
+    client, connection, transport, protocol = api_client
+    notifies = []
+
+    def on_bluetooth_gatt_notify(handle: int, data: bytearray) -> None:
+        notifies.append((handle, data))
+
+    notify_task = asyncio.create_task(
+        client.bluetooth_gatt_start_notify(1234, 1, on_bluetooth_gatt_notify)
+    )
+    await asyncio.sleep(0)
+    notify_response: message.Message = BluetoothGATTNotifyResponse(
+        address=1234, handle=1
+    )
+    data_response: message.Message = BluetoothGATTNotifyDataResponse(
+        address=1234, handle=1, data=b"gotit"
+    )
+    protocol.data_received(
+        generate_plaintext_packet(notify_response)
+        + generate_plaintext_packet(data_response)
+    )
+
+    await notify_task
+    assert notifies == [(1, b"gotit")]
+
+    second_data_response: message.Message = BluetoothGATTNotifyDataResponse(
+        address=1234, handle=1, data=b"after finished"
+    )
+    protocol.data_received(generate_plaintext_packet(second_data_response))
+    assert notifies == [(1, b"gotit"), (1, b"after finished")]
+
+
 @pytest.mark.asyncio
 async def test_subscribe_logs(auth_client: APIClient) -> None:
     send = patch_response_callback(auth_client)