diff --git a/src/quartz_api/internal/backends/test_utils.py b/src/quartz_api/internal/backends/test_utils.py new file mode 100644 index 00000000..c7525c15 --- /dev/null +++ b/src/quartz_api/internal/backends/test_utils.py @@ -0,0 +1,20 @@ +import datetime as dt +import unittest + +from .utils import get_window + + +class TestGetWindow(unittest.TestCase): + def test_get_window_with_reference_time(self) -> None: + reference_time = dt.datetime(2024, 1, 10, 15, 30, tzinfo=dt.UTC) + + start, end = get_window(reference_time=reference_time) + + self.assertEqual( + start, + dt.datetime(2024, 1, 8, 0, 0, tzinfo=dt.UTC), + ) + self.assertEqual( + end, + dt.datetime(2024, 1, 12, 0, 0, tzinfo=dt.UTC), + ) diff --git a/src/quartz_api/internal/backends/utils.py b/src/quartz_api/internal/backends/utils.py index 98ec4566..9fa00254 100644 --- a/src/quartz_api/internal/backends/utils.py +++ b/src/quartz_api/internal/backends/utils.py @@ -3,17 +3,27 @@ import datetime as dt -def get_window() -> tuple[dt.datetime, dt.datetime]: - """Returns the start and end of the window for timeseries data.""" +def get_window( + reference_time: dt.datetime | None = None, +) -> tuple[dt.datetime, dt.datetime]: + """Returns the start and end of the window for timeseries data. + + Args: + reference_time: The time to base the window on. If None, defaults to now (UTC). + """ + if reference_time is None: + reference_time = dt.datetime.now(tz=dt.UTC) + # Window start is the beginning of the day two days ago - start = (dt.datetime.now(tz=dt.UTC) - dt.timedelta(days=2)).replace( + start = (reference_time - dt.timedelta(days=2)).replace( hour=0, minute=0, second=0, microsecond=0, ) + # Window end is the beginning of the day two days ahead - end = (dt.datetime.now(tz=dt.UTC) + dt.timedelta(days=2)).replace( + end = (reference_time + dt.timedelta(days=2)).replace( hour=0, minute=0, second=0,