From 3389e3e1edbd94cdf06c13b1a967095e16858670 Mon Sep 17 00:00:00 2001
From: Alan Rosenthal <alanrosenthal@google.com>
Date: Thu, 11 Jul 2024 11:01:07 -0400
Subject: [PATCH] device.py: Update get_characteristics_by_uuid()

`get_characteristics_by_uuid()` now allows a UUID to be passed to the
service param. This allows for users to easily query for a service uuid
and characteristic uuid with one API.
---
 bumble/device.py   | 15 +++++++++++++-
 tests/gatt_test.py | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/bumble/device.py b/bumble/device.py
index 563bbf87..dcabf765 100644
--- a/bumble/device.py
+++ b/bumble/device.py
@@ -27,6 +27,7 @@
 from dataclasses import dataclass, field
 from enum import Enum, IntEnum
 import functools
+import itertools
 import json
 import logging
 import secrets
@@ -1172,8 +1173,20 @@ def get_services_by_uuid(self, uuid: core.UUID) -> List[gatt_client.ServiceProxy
         return self.gatt_client.get_services_by_uuid(uuid)
 
     def get_characteristics_by_uuid(
-        self, uuid: core.UUID, service: Optional[gatt_client.ServiceProxy] = None
+        self,
+        uuid: core.UUID,
+        service: Optional[Union[gatt_client.ServiceProxy, core.UUID]] = None,
     ) -> List[gatt_client.CharacteristicProxy]:
+        if isinstance(service, core.UUID):
+            return list(
+                itertools.chain(
+                    *[
+                        self.get_characteristics_by_uuid(uuid, s)
+                        for s in self.get_services_by_uuid(service)
+                    ]
+                )
+            )
+
         return self.gatt_client.get_characteristics_by_uuid(uuid, service)
 
     def create_service_proxy(self, proxy_class: Type[_PROXY_CLASS]) -> _PROXY_CLASS:
diff --git a/tests/gatt_test.py b/tests/gatt_test.py
index e3c92097..9d8454a8 100644
--- a/tests/gatt_test.py
+++ b/tests/gatt_test.py
@@ -1146,6 +1146,56 @@ def test_get_attribute_group():
     )
 
 
+# -----------------------------------------------------------------------------
+@pytest.mark.asyncio
+async def test_get_characteristics_by_uuid():
+    [client, server] = LinkedDevices().devices[:2]
+
+    characteristic1 = Characteristic(
+        '1234',
+        Characteristic.Properties.READ | Characteristic.Properties.NOTIFY,
+        Characteristic.READABLE,
+        bytes([1, 2, 3]),
+    )
+    characteristic2 = Characteristic(
+        '5678',
+        Characteristic.Properties.READ | Characteristic.Properties.NOTIFY,
+        Characteristic.READABLE,
+        bytes([1, 2, 3]),
+    )
+    service1 = Service(
+        'ABCD',
+        [characteristic1, characteristic2],
+    )
+    service2 = Service(
+        'FFFF',
+        [characteristic1],
+    )
+
+    server.add_services([service1, service2])
+
+    await client.power_on()
+    await server.power_on()
+    connection = await client.connect(server.random_address)
+    peer = Peer(connection)
+
+    await peer.discover_services()
+    await peer.discover_characteristics()
+    c = peer.get_characteristics_by_uuid(uuid=UUID('1234'))
+    assert len(c) == 2
+    assert isinstance(c[0], CharacteristicProxy)
+    c = peer.get_characteristics_by_uuid(uuid=UUID('1234'), service=UUID('ABCD'))
+    assert len(c) == 1
+    assert isinstance(c[0], CharacteristicProxy)
+    c = peer.get_characteristics_by_uuid(uuid=UUID('1234'), service=UUID('AAAA'))
+    assert len(c) == 0
+
+    s = peer.get_services_by_uuid(uuid=UUID('ABCD'))
+    assert len(s) == 1
+    c = peer.get_characteristics_by_uuid(uuid=UUID('1234'), service=s[0])
+    assert len(s) == 1
+
+
 # -----------------------------------------------------------------------------
 if __name__ == '__main__':
     logging.basicConfig(level=os.environ.get('BUMBLE_LOGLEVEL', 'INFO').upper())