[MovingWindow] Accept a time range in the count_valid method

It retains the original behaviour of counting all the valid samples in
the buffer when no time range is specified.

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

Author: shsms

src/frequenz/sdk/timeseries/_moving_window.py

@@ -355,14 +355,24 @@ async def sink_buffer(sample: Sample[Quantity]) -> None:
             asyncio.create_task(self._resampler.resample(), name="resample")
         )
 
-    def count_valid(self) -> int:
-        """
-        Count the number of valid samples in this `MovingWindow`.
+    def count_valid(
+        self, *, since: datetime | None = None, until: datetime | None = None
+    ) -> int:
+        """Count the number of valid samples in this `MovingWindow`.
+
+        If `since` and `until` are provided, the count is limited to the samples between
+        (and including) the given timestamps.
+
+        Args:
+            since: The timestamp from which to start counting.  If `None`, the oldest
+                timestamp of the buffer is used.
+            until: The timestamp until (and including) which to count.  If `None`, the
+                newest timestamp of the buffer is used.
 
         Returns:
             The number of valid samples in this `MovingWindow`.
         """
-        return self._buffer.count_valid()
+        return self._buffer.count_valid(since=since, until=until)
 
     def count_covered(self) -> int:
         """Count the number of samples that are covered by the oldest and newest valid samples.

src/frequenz/sdk/timeseries/_ringbuffer/buffer.py

@@ -678,31 +678,48 @@ def count_covered(self) -> int:
             // self._sampling_period.total_seconds()
         )
 
-    def count_valid(self) -> int:
-        """Count the number of valid items that this buffer currently holds.
+    def count_valid(
+        self, *, since: datetime | None = None, until: datetime | None = None
+    ) -> int:
+        """Count the number of valid items in this buffer.
+
+        If `since` and `until` are provided, the count is limited to the items between
+        (and including) the given timestamps.
+
+        Args:
+            since: The timestamp from which to start counting.  If `None`, the oldest
+                timestamp in the buffer is used.
+            until: The timestamp until (and including) which to count.  If `None`, the
+                newest timestamp in the buffer is used.
 
         Returns:
             The number of valid items in this buffer.
         """
-        if self._timestamp_newest == self._TIMESTAMP_MIN:
+        if since is None or since < self._timestamp_oldest:
+            since = self._timestamp_oldest
+        if until is None or until > self._timestamp_newest:
+            until = self._timestamp_newest
+
+        if until == self._TIMESTAMP_MIN or until < since:
             return 0
 
         # Sum of all elements in the gap ranges
         sum_missing_entries = max(
             0,
             sum(
                 (
-                    gap.end
+                    min(gap.end, until + self._sampling_period)
                     # Don't look further back than oldest timestamp
-                    - max(gap.start, self._timestamp_oldest)
+                    - max(gap.start, since)
                 )
                 // self._sampling_period
                 for gap in self._gaps
+                if gap.start <= until and gap.end >= since
             ),
         )
 
-        start_pos = self.to_internal_index(self._timestamp_oldest)
-        end_pos = self.to_internal_index(self._timestamp_newest)
+        start_pos = self.to_internal_index(since)
+        end_pos = self.to_internal_index(until)
 
         if end_pos < start_pos:
             return len(self._buffer) - start_pos + end_pos + 1 - sum_missing_entries

tests/timeseries/test_moving_window.py

@@ -212,19 +212,91 @@ async def test_access_empty_window() -> None:
 
 async def test_window_size() -> None:
     """Test the size of the window."""
-    window, sender = init_moving_window(timedelta(seconds=5))
+    window, sender = init_moving_window(timedelta(seconds=10))
     async with window:
-        assert window.capacity == 5, "Wrong window capacity"
+        assert window.capacity == 10, "Wrong window capacity"
         assert window.count_valid() == 0, "Window should be empty"
         assert window.count_covered() == 0, "Window should be empty"
+
         await push_logical_meter_data(sender, range(0, 2))
-        assert window.capacity == 5, "Wrong window capacity"
+        assert window.capacity == 10, "Wrong window capacity"
         assert window.count_valid() == 2, "Window should be partially full"
         assert window.count_covered() == 2, "Window should be partially full"
-        await push_logical_meter_data(sender, range(2, 20))
-        assert window.capacity == 5, "Wrong window capacity"
-        assert window.count_valid() == 5, "Window should be full"
-        assert window.count_covered() == 5, "Window should be full"
+
+        newest_ts = window.newest_timestamp
+        assert newest_ts is not None and newest_ts == UNIX_EPOCH + timedelta(seconds=1)
+
+        await push_logical_meter_data(sender, range(2, 5), start_ts=newest_ts)
+        assert window.capacity == 10, "Wrong window capacity"
+        assert window.count_valid() == 4, "Window should be partially full"
+        assert window.count_covered() == 4, "Window should be partially full"
+
+        newest_ts = window.newest_timestamp
+        assert newest_ts is not None and newest_ts == UNIX_EPOCH + timedelta(seconds=3)
+
+        await push_logical_meter_data(sender, range(5, 12), start_ts=newest_ts)
+        assert window.capacity == 10, "Wrong window capacity"
+        assert window.count_valid() == 10, "Window should be full"
+        assert window.count_covered() == 10, "Window should be full"
+
+        assert window.count_valid(since=UNIX_EPOCH + timedelta(seconds=1)) == 9
+        assert window.count_valid(until=UNIX_EPOCH + timedelta(seconds=2)) == 3
+        assert (
+            window.count_valid(
+                since=UNIX_EPOCH + timedelta(seconds=1),
+                until=UNIX_EPOCH + timedelta(seconds=1),
+            )
+            == 1
+        )
+        assert (
+            window.count_valid(
+                since=UNIX_EPOCH + timedelta(seconds=3),
+                until=UNIX_EPOCH + timedelta(seconds=8),
+            )
+            == 6
+        )
+        assert (
+            window.count_valid(
+                since=UNIX_EPOCH + timedelta(seconds=8),
+                until=UNIX_EPOCH + timedelta(seconds=3),
+            )
+            == 0
+        )
+
+        newest_ts = window.newest_timestamp
+        assert newest_ts is not None and newest_ts == UNIX_EPOCH + timedelta(seconds=9)
+        assert window.oldest_timestamp == UNIX_EPOCH
+
+        await push_logical_meter_data(sender, range(5, 12), start_ts=newest_ts)
+        assert window.capacity == 10, "Wrong window capacity"
+        assert window.count_valid() == 10, "Window should be full"
+        assert window.count_covered() == 10, "Window should be full"
+
+        newest_ts = window.newest_timestamp
+        assert newest_ts is not None and newest_ts == UNIX_EPOCH + timedelta(seconds=15)
+        assert window.oldest_timestamp == UNIX_EPOCH + timedelta(seconds=6)
+
+        assert (
+            window.count_valid(
+                since=UNIX_EPOCH + timedelta(seconds=1),
+                until=UNIX_EPOCH + timedelta(seconds=5),
+            )
+            == 0
+        )
+        assert (
+            window.count_valid(
+                since=UNIX_EPOCH + timedelta(seconds=3),
+                until=UNIX_EPOCH + timedelta(seconds=8),
+            )
+            == 3
+        )
+        assert (
+            window.count_valid(
+                since=UNIX_EPOCH + timedelta(seconds=6),
+                until=UNIX_EPOCH + timedelta(seconds=20),
+            )
+            == 10
+        )
 
 
 # pylint: disable=redefined-outer-name