[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,44 @@ 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.
+
+        This function returns after `n` new samples are available in the MovingWindow,
+        without considering whether the new samples are valid.  The validity of the
+        samples can be verified by calling the
+        [`count_valid`][frequenz.sdk.timeseries.MovingWindow.count_valid] method.
+
+        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:
+            return
+        if n < 0:
+            raise ValueError("The number of samples to wait for must be 0 or greater.")
+        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 = (
+            # Start from the next expected timestamp.
+            self.newest_timestamp + self.sampling_period
+            if self.newest_timestamp is not None
+            else None
+        )
+        while True:
+            async with self._condition_new_sample:
+                # Every time a new sample is received, this condition gets notified and
+                # will wake up.
+                _ = await self._condition_new_sample.wait()
+            if self.count_covered(since=start_timestamp) >= n:
+                return
+
     async def _run_impl(self) -> None:
         """Awaits samples from the receiver and updates the underlying ring buffer.
 
@@ -331,6 +371,9 @@ async def _run_impl(self) -> None:
                     await self._resampler_sender.send(sample)
                 else:
                     self._buffer.update(sample)
+                    async with self._condition_new_sample:
+                        # Wake up all coroutines waiting for new samples.
+                        self._condition_new_sample.notify_all()
 
         except asyncio.CancelledError:
             _logger.info("MovingWindow task has been cancelled.")
@@ -344,6 +387,9 @@ def _configure_resampler(self) -> None:
 
         async def sink_buffer(sample: Sample[Quantity]) -> None:
             self._buffer.update(sample)
+            async with self._condition_new_sample:
+                # Wake up all coroutines waiting for new samples.
+                self._condition_new_sample.notify_all()
 
         resampler_channel = Broadcast[Sample[Quantity]](name="average")
         self._resampler_sender = resampler_channel.new_sender()

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
 
@@ -29,6 +30,7 @@ async def push_logical_meter_data(
     sender: Sender[Sample[Quantity]],
     test_seq: Sequence[float | None],
     start_ts: datetime = UNIX_EPOCH,
+    fake_time: time_machine.Coordinates | None = None,
 ) -> None:
     """Push data in the passed sender to mock `LogicalMeter` behaviour.
 
@@ -38,23 +40,29 @@ async def push_logical_meter_data(
         sender: Sender for pushing resampled samples to the `MovingWindow`.
         test_seq: The Sequence that is pushed into the `MovingWindow`.
         start_ts: The start timestamp of the `MovingWindow`.
+        fake_time: The fake time object to shift the time.
     """
     for i, j in zip(test_seq, range(0, len(test_seq))):
         timestamp = start_ts + timedelta(seconds=j)
         await sender.send(
             Sample(timestamp, Quantity(float(i)) if i is not None else None)
         )
+        if fake_time is not None:
+            await asyncio.sleep(1.0)
+            fake_time.shift(1)
 
     await asyncio.sleep(0.0)
 
 
 def init_moving_window(
     size: timedelta,
+    resampler_config: ResamplerConfig | None = None,
 ) -> tuple[MovingWindow, Sender[Sample[Quantity]]]:
     """Initialize the moving window with given shape.
 
     Args:
         size: The size of the `MovingWindow`
+        resampler_config: The resampler configuration.
 
     Returns:
         tuple[MovingWindow, Sender[Sample]]: A pair of sender and `MovingWindow`.
@@ -65,6 +73,7 @@ def init_moving_window(
         size=size,
         resampled_data_recv=lm_chan.new_receiver(),
         input_sampling_period=timedelta(seconds=1),
+        resampler_config=resampler_config,
     )
     return window, lm_tx
 
@@ -363,6 +372,149 @@ def assert_valid_and_covered_counts(
         )
 
 
+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)
+        # After pushing 5 values, the `wait_for_samples` task should be done.
+        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)
+        # The task should not be done yet, since we have only pushed 4 values.
+        assert not task.done()
+
+        await push_logical_meter_data(
+            sender, [1], start_ts=UNIX_EPOCH + timedelta(seconds=9)
+        )
+        await asyncio.sleep(0)
+        # After pushing the last value, the task should be done.
+        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 0 or greater."),
+        ):
+            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()
+
+        task = asyncio.create_task(window.wait_for_samples(5))
+        await asyncio.sleep(0)
+        await push_logical_meter_data(
+            sender, [1, 2, None, 4, None], start_ts=UNIX_EPOCH + timedelta(seconds=30)
+        )
+        await asyncio.sleep(0)
+        # `None` values *are* counted towards the number of samples to wait for.
+        assert task.done()
+
+
+async def test_wait_for_samples_with_resampling(
+    fake_time: time_machine.Coordinates,
+) -> None:
+    """Test waiting for samples in a moving window with resampling."""
+    window, sender = init_moving_window(
+        timedelta(seconds=20), ResamplerConfig(resampling_period=timedelta(seconds=2))
+    )
+    async with window:
+        task = asyncio.create_task(window.wait_for_samples(3))
+        await asyncio.sleep(0)
+        assert not task.done()
+        await push_logical_meter_data(sender, range(0, 7), fake_time=fake_time)
+        assert task.done()
+
+        task = asyncio.create_task(window.wait_for_samples(10))
+        await push_logical_meter_data(
+            sender,
+            range(0, 11),
+            fake_time=fake_time,
+            start_ts=UNIX_EPOCH + timedelta(seconds=7),
+        )
+        assert window.count_covered() == 8
+        assert not task.done()
+
+        await push_logical_meter_data(
+            sender,
+            range(0, 5),
+            fake_time=fake_time,
+            start_ts=UNIX_EPOCH + timedelta(seconds=18),
+        )
+        assert window.count_covered() == 10
+        assert not task.done()
+
+        await push_logical_meter_data(
+            sender,
+            range(0, 6),
+            fake_time=fake_time,
+            start_ts=UNIX_EPOCH + timedelta(seconds=23),
+        )
+        assert window.count_covered() == 10
+        assert window.count_valid() == 10
+        assert task.done()
+
+        task = asyncio.create_task(window.wait_for_samples(5))
+        await push_logical_meter_data(
+            sender,
+            [1, 2, None, None, None, None, None, None, None, None],
+            fake_time=fake_time,
+            start_ts=UNIX_EPOCH + timedelta(seconds=29),
+        )
+        assert window.count_covered() == 10
+        assert window.count_valid() == 8
+        assert task.done()
+
+        task = asyncio.create_task(window.wait_for_samples(5))
+        await push_logical_meter_data(
+            sender,
+            [None, 4, None, None, None, None, None, None, None, 5],
+            fake_time=fake_time,
+            start_ts=UNIX_EPOCH + timedelta(seconds=39),
+        )
+        assert window.count_covered() == 10
+        assert window.count_valid() == 7
+        assert task.done()
+
+
 # pylint: disable=redefined-outer-name
 async def test_resampling_window(fake_time: time_machine.Coordinates) -> None:
     """Test resampling in MovingWindow."""