diff --git a/src/frequenz/sdk/timeseries/_moving_window.py b/src/frequenz/sdk/timeseries/_moving_window.py index 5be8a5e34..d180e1267 100644 --- a/src/frequenz/sdk/timeseries/_moving_window.py +++ b/src/frequenz/sdk/timeseries/_moving_window.py @@ -138,6 +138,13 @@ def __init__( # pylint: disable=too-many-arguments self._resampler_sender: Sender[Sample] | None = None self._resampler_task: asyncio.Task[None] | None = None + self._wait_for_num_samples: int = 0 + """The number of samples to wait for before the wait_for_num_samples channels + sends out an event.""" + self._wait_for_samples_channel = Broadcast[None]( + "Wait for number of samples channel." + ) + if resampler_config: assert ( resampler_config.resampling_period <= size @@ -169,6 +176,9 @@ async def _run_impl(self) -> None: Raises: asyncio.CancelledError: if the MovingWindow task is cancelled. """ + received_samples_count = 0 + wait_for_samples_sender = self._wait_for_samples_channel.new_sender() + try: async for sample in self._resampled_data_recv: _logger.debug("Received new sample: %s", sample) @@ -177,12 +187,48 @@ async def _run_impl(self) -> None: else: self._buffer.update(sample) + # count the number of samples and send out a trigger when it matches + # the number of samples to wait for. + received_samples_count += 1 + if self._wait_for_num_samples != 0: + if received_samples_count == self._wait_for_num_samples: + received_samples_count = 0 + await wait_for_samples_sender.send(None) + except asyncio.CancelledError: _logger.info("MovingWindow task has been cancelled.") raise _logger.error("Channel has been closed") + def set_sample_counter(self, num_samples: int) -> None: + """Set the number of samples to wait for until the sample counter triggers. + + Args: + num_samples: The number of samples to wait for. + + Raises: + ValueError: if the number of samples is less than or equal to zero. + """ + if num_samples <= 0: + raise ValueError( + "The number of samples to wait for should be greater than zero." + ) + self._wait_for_num_samples = num_samples + + def new_sample_count_receiver(self) -> Receiver[None]: + """Wait until a given number of samples has been received. + + The sample counter is updated irrespective of whether this + method is called or not. Thus this might trigger when + a smaller number of samples than the given number has been + updated. + + Returns: + A receiver that triggers after a number of samples arrived. + """ + return self._wait_for_samples_channel.new_receiver() + async def stop(self) -> None: """Cancel the running tasks and stop the MovingWindow.""" await cancel_and_await(self._update_window_task) diff --git a/tests/timeseries/test_moving_window.py b/tests/timeseries/test_moving_window.py index f692752b3..e4de7f3ef 100644 --- a/tests/timeseries/test_moving_window.py +++ b/tests/timeseries/test_moving_window.py @@ -104,6 +104,27 @@ async def test_window_size() -> None: assert len(window) == 5 +@pytest.mark.parametrize("samples_to_wait_for", [-1, 0, 1, 10]) +async def test_wait_for_samples(samples_to_wait_for: int) -> None: + """Test waiting for samples.""" + window, sender = init_moving_window(timedelta(seconds=1)) + + if samples_to_wait_for <= 0: + with pytest.raises(ValueError): + window.set_sample_counter(samples_to_wait_for) + return + sample_count_recv = window.new_sample_count_receiver() + + window.set_sample_counter(samples_to_wait_for) + + # asyncio.create_task(push_data_delayed()) + for i in range(0, samples_to_wait_for): + await sender.send( + Sample(datetime.now(tz=timezone.utc) + timedelta(seconds=i), 1.0) + ) + await sample_count_recv.receive() + + # pylint: disable=redefined-outer-name async def test_resampling_window(fake_time: time_machine.Coordinates) -> None: """Test resampling in MovingWindow."""