[MovingWindow] Add a wait_for_samples method

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/sdk/timeseries/_moving_window.py

@@ -190,6 +190,8 @@ def __init__(  # pylint: disable=too-many-arguments
             align_to=align_to,
         )
 
+        self._condition_new_sample = asyncio.Condition()
+
     def start(self) -> None:
         """Start the MovingWindow.
 
@@ -318,6 +320,34 @@ def window(
             start, end, force_copy=force_copy, fill_value=fill_value
         )
 
+    async def wait_for_samples(self, n: int) -> None:
+        """Wait until the next `n` samples are available in the MovingWindow.
+
+        Args:
+            n: The number of samples to wait for.
+
+        Raises:
+            ValueError: If `n` is less than or equal to 0 or greater than the capacity
+                of the MovingWindow.
+        """
+        if n <= 0:
+            raise ValueError(
+                "The number of samples to wait for must be greater than 0."
+            )
+        if n > self.capacity:
+            raise ValueError(
+                "The number of samples to wait for must be less than or equal to the "
+                + f"capacity of the MovingWindow ({self.capacity})."
+            )
+        start_timestamp = self.newest_timestamp
+        if n < self.capacity:
+            n += self.count_valid(since=start_timestamp)
+        while True:
+            async with self._condition_new_sample:
+                _ = await self._condition_new_sample.wait()
+            if self.count_valid(since=start_timestamp) >= n:
+                return
+
     async def _run_impl(self) -> None:
         """Awaits samples from the receiver and updates the underlying ring buffer.
 
@@ -330,6 +360,8 @@ async def _run_impl(self) -> None:
                 if self._resampler and self._resampler_sender:
                     await self._resampler_sender.send(sample)
                 else:
+                    async with self._condition_new_sample:
+                        self._condition_new_sample.notify_all()
                     self._buffer.update(sample)
 
         except asyncio.CancelledError:
@@ -344,6 +376,8 @@ def _configure_resampler(self) -> None:
 
         async def sink_buffer(sample: Sample[Quantity]) -> None:
             if sample.value is not None:
+                async with self._condition_new_sample:
+                    self._condition_new_sample.notify_all()
                 self._buffer.update(sample)
 
         resampler_channel = Broadcast[Sample[Quantity]](name="average")

tests/timeseries/test_moving_window.py

@@ -4,6 +4,7 @@
 """Tests for the moving window."""
 
 import asyncio
+import re
 from collections.abc import Sequence
 from datetime import datetime, timedelta, timezone
 
@@ -299,6 +300,73 @@ async def test_window_size() -> None:
         )
 
 
+async def test_wait_for_samples() -> None:
+    """Test waiting for samples in the window."""
+    window, sender = init_moving_window(timedelta(seconds=10))
+    async with window:
+        task = asyncio.create_task(window.wait_for_samples(5))
+        await asyncio.sleep(0)
+        assert not task.done()
+        await push_logical_meter_data(sender, range(0, 5))
+        await asyncio.sleep(0)
+        assert task.done()
+
+        task = asyncio.create_task(window.wait_for_samples(5))
+        await asyncio.sleep(0)
+        await push_logical_meter_data(
+            sender, [1, 2, 3, 4], start_ts=UNIX_EPOCH + timedelta(seconds=5)
+        )
+        await asyncio.sleep(0)
+        assert not task.done()
+
+        await push_logical_meter_data(
+            sender, [1], start_ts=UNIX_EPOCH + timedelta(seconds=9)
+        )
+        await asyncio.sleep(0)
+        assert task.done()
+
+        task = asyncio.create_task(window.wait_for_samples(-1))
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "The number of samples to wait for must be greater than 0."
+            ),
+        ):
+            await task
+
+        task = asyncio.create_task(window.wait_for_samples(20))
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "The number of samples to wait for must be less than or equal to the "
+                + "capacity of the MovingWindow (10)."
+            ),
+        ):
+            await task
+
+        task = asyncio.create_task(window.wait_for_samples(4))
+        await asyncio.sleep(0)
+        await push_logical_meter_data(
+            sender, range(0, 10), start_ts=UNIX_EPOCH + timedelta(seconds=10)
+        )
+        await asyncio.sleep(0)
+        assert task.done()
+
+        task = asyncio.create_task(window.wait_for_samples(10))
+        await asyncio.sleep(0)
+        await push_logical_meter_data(
+            sender, range(0, 5), start_ts=UNIX_EPOCH + timedelta(seconds=20)
+        )
+        await asyncio.sleep(0)
+        assert not task.done()
+
+        await push_logical_meter_data(
+            sender, range(10, 15), start_ts=UNIX_EPOCH + timedelta(seconds=25)
+        )
+        await asyncio.sleep(0)
+        assert task.done()
+
+
 # pylint: disable=redefined-outer-name
 async def test_resampling_window(fake_time: time_machine.Coordinates) -> None:
     """Test resampling in MovingWindow."""