1import numpy as np2import pandas as pd34from chronos_x.data.calendar import sessions, slot_index, REGULAR_SLOTS, EARLY_CLOSE_SLOTS567def test_regular_and_early_close_slots():8 s = {x.date.date().isoformat(): x for x in sessions("2024-11-25", "2024-12-02")}9 assert s["2024-11-29"].n_slots == EARLY_CLOSE_SLOTS and s["2024-11-29"].early_close10 assert s["2024-11-27"].n_slots == REGULAR_SLOTS and not s["2024-11-27"].early_close11 assert "2024-11-28" not in s # Thanksgiving121314def test_dst_open_in_utc():15 s = {x.date.date().isoformat(): x for x in sessions("2024-03-07", "2024-03-12")}16 assert s["2024-03-08"].open_utc == pd.Timestamp("2024-03-08T14:30:00Z") # EST17 assert s["2024-03-11"].open_utc == pd.Timestamp("2024-03-11T13:30:00Z") # EDT18 assert s["2024-03-11"].close_utc == pd.Timestamp("2024-03-11T20:00:00Z")192021def test_slot_index_bar_open_convention():22 s = sessions("2024-03-11", "2024-03-11")[0]23 ts = np.array(["2024-03-11T13:30:00", "2024-03-11T13:31:00", "2024-03-11T20:00:00", "2024-03-11T20:01:00",24 "2024-03-11T13:29:00", "2024-03-11T13:30:30"], dtype="datetime64[s]")25 k = slot_index(ts, np.full(len(ts), s.open_utc.to_datetime64().astype("datetime64[s]")), np.full(len(ts), s.n_slots))26 assert k.tolist() == [0, 1, 390, -1, -1, -1]27