From 5264b902ba0d9b37dc7d4168f3b3e070d88ca6e9 Mon Sep 17 00:00:00 2001
From: ttu <tomi.tuhkanen@iki.fi>
Date: Thu, 23 Jan 2025 21:41:09 +0200
Subject: [PATCH] Add maximum number of items to fetch

---
 ruuvitag_sensor/adapters/bleak_ble.py | 13 ++++++++++---
 ruuvitag_sensor/ruuvi.py              | 19 +++++++++++--------
 2 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/ruuvitag_sensor/adapters/bleak_ble.py b/ruuvitag_sensor/adapters/bleak_ble.py
index e873dfc..7c64c0e 100644
--- a/ruuvitag_sensor/adapters/bleak_ble.py
+++ b/ruuvitag_sensor/adapters/bleak_ble.py
@@ -126,13 +126,16 @@ async def get_first_data(mac: str, bt_device: str = "") -> RawData:
 
         return data or ""
 
-    async def get_history_data(self, mac: str, start_time: Optional[datetime] = None) -> List[dict]:
+    async def get_history_data(
+        self, mac: str, start_time: Optional[datetime] = None, max_items: Optional[int] = None
+    ) -> List[dict]:
         """
         Get history data from a RuuviTag using GATT connection.
 
         Args:
             mac (str): MAC address of the RuuviTag
             start_time (datetime, optional): Start time for history data
+            max_items (int, optional): Maximum number of history entries to fetch
 
         Returns:
             List[dict]: List of historical sensor readings
@@ -184,13 +187,17 @@ def notification_handler(_, data: bytearray):
             log.debug("Requested history data from device %s", mac)
 
             # Wait for initial notification
-            await asyncio.wait_for(notification_received.wait(), timeout=5.0)
+            await asyncio.wait_for(notification_received.wait(), timeout=10.0)
 
             # Wait for more data
             try:
                 while True:
                     notification_received.clear()
-                    await asyncio.wait_for(notification_received.wait(), timeout=1.0)
+                    await asyncio.wait_for(notification_received.wait(), timeout=5.0)
+                    # Check if we've reached the maximum number of items
+                    if max_items and len(history_data) >= max_items:
+                        log.debug("Reached maximum number of items (%d)", max_items)
+                        break
             except asyncio.TimeoutError:
                 # No more data received for 1 second - assume transfer complete
                 pass
diff --git a/ruuvitag_sensor/ruuvi.py b/ruuvitag_sensor/ruuvi.py
index 5f961d0..c598a67 100644
--- a/ruuvitag_sensor/ruuvi.py
+++ b/ruuvitag_sensor/ruuvi.py
@@ -346,29 +346,35 @@ def _parse_data(
         return (mac_to_send, decoded)
 
     @staticmethod
-    async def get_history_async(mac: str, start_time: Optional[datetime] = None) -> List[SensorData]:
+    async def get_history_async(
+        mac: str, start_time: Optional[datetime] = None, max_items: Optional[int] = None
+    ) -> List[SensorData]:
         """
         Get history data from a RuuviTag that supports it (firmware 3.30.0+)
 
         Args:
             mac (str): MAC address of the RuuviTag
             start_time (datetime, optional): Start time for history data. If None, gets all available data
+            max_items (int, optional): Maximum number of history entries to fetch. If None, gets all available data
 
         Returns:
             List[SensorData]: List of historical sensor readings
         """
         throw_if_not_async_adapter(ble)
-        return await ble.get_history_data(mac, start_time)
+        return await ble.get_history_data(mac, start_time, max_items)
 
     @staticmethod
-    async def download_history(mac: str, start_time: Optional[datetime] = None, timeout: int = 300) -> List[SensorData]:
+    async def download_history(
+        mac: str, start_time: Optional[datetime] = None, timeout: int = 300, max_items: Optional[int] = None
+    ) -> List[SensorData]:
         """
         Download history data from a RuuviTag. Requires firmware version 3.30.0 or newer.
 
         Args:
             mac (str): MAC address of the RuuviTag. On macOS use UUID instead.
             start_time (Optional[datetime]): If provided, only get data from this time onwards
-            timeout (int): Maximum time in seconds to wait for history download (default: 30)
+            timeout (int): Maximum time in seconds to wait for history download (default: 300)
+            max_items (Optional[int]): Maximum number of history entries to fetch. If None, gets all available data
 
         Returns:
             List[SensorData]: List of historical measurements, ordered by timestamp
@@ -379,11 +385,8 @@ async def download_history(mac: str, start_time: Optional[datetime] = None, time
         """
         throw_if_not_async_adapter(ble)
 
-        # if not re.match("[0-9A-F]{2}(:[0-9A-F]{2}){5}$", mac.upper()):
-        #     raise ValueError(f"Invalid MAC address: {mac}")
-
         try:
-            history = await asyncio.wait_for(ble.get_history_data(mac, start_time), timeout=timeout)
+            history = await asyncio.wait_for(ble.get_history_data(mac, start_time, max_items), timeout=timeout)
 
             # Sort by timestamp if present
             if history and "timestamp" in history[0]: