feature: also test DatetimeIndex against DST transitions

Signed-off-by: F.N. Claessen <felix@seita.nl>
SeitaBV · Mar 23, 2024 · 14d3ebb · 14d3ebb
1 parent 1ce00fb
commit 14d3ebb
Showing 1 changed file with 21 additions and 3 deletions.
diff --git a/timely_beliefs/sensors/func_store/test_knowledge_horizons.py b/timely_beliefs/sensors/func_store/test_knowledge_horizons.py
@@ -147,24 +147,42 @@ def test_x_years_ago_at_date_with_dst():
     - 2024-03-30
     """
 
-    knowledge_func_params_dst = dict(
+    knowledge_func_params = dict(
         x=1, month=3, day=28, z="Europe/Amsterdam"
     )  # before first DST transition 2024
     assert x_years_ago_at_date(
         event_start=timezone("Europe/Amsterdam").localize(datetime(2024, 3, 28, 0)),
-        **knowledge_func_params_dst,
+        **knowledge_func_params,
     ) == timedelta(
         days=366, hours=1
     )  # 365 + 1 day (because of the leap day on 2024-02-29) + 1 hour (fall transition)
 
     # Try 4 days later, at which we crossed the spring DST transition
     assert x_years_ago_at_date(
         event_start=timezone("Europe/Amsterdam").localize(datetime(2024, 4, 1, 0)),
-        **knowledge_func_params_dst,
+        **knowledge_func_params,
     ) == timedelta(
         days=370, hours=0
     )  # 0 hours (fall and spring transitions cancelled each other out)
 
+    # Repeat test with pd.DatetimeIndex instead
+    event_start = pd.DatetimeIndex(
+        [
+            "2024-03-28T00:00:00",
+            "2024-04-01T00:00:00",
+        ],
+        tz="Europe/Amsterdam",
+    )
+    assert_index_equal(
+        x_years_ago_at_date(event_start=event_start, **knowledge_func_params),
+        pd.TimedeltaIndex(
+            [
+                timedelta(days=366, hours=1),
+                timedelta(days=370, hours=0),
+            ]
+        ),
+    )
+
 
 @pytest.mark.parametrize(
     "event_start",