from datetime import datetime

import pendulum
from dagster.core.definitions.time_window_partitions import TimeWindow, daily_partitioned_config

DATE_FORMAT = "%Y-%m-%d"


def time_window(start: str, end: str) -> TimeWindow:
    return TimeWindow(pendulum.parse(start), pendulum.parse(end))


def test_daily_partitions():
    @daily_partitioned_config(start_date="2021-05-05")
    def my_partitioned_config(_start, _end):
        return {}

    assert [
        partition.value
        for partition in my_partitioned_config.partitions_def.get_partitions(
            datetime.strptime("2021-05-07", DATE_FORMAT)
        )
    ] == [
        time_window("2021-05-05", "2021-05-06"),
        time_window("2021-05-06", "2021-05-07"),
    ]
