from datetime import datetime
from zfs_smart_snapshot.main import keep_year, keep_month, keep_day, keep_hour, keep_frequent


dates = [datetime(2018, 2, 3), datetime(2019, 3, 4), datetime(2020, 4, 5), datetime(2021, 1, 2),
         datetime(2021, 2, 2), datetime(2021, 2, 3), datetime(2021, 3, 3), datetime(2021, 3, 4),
         datetime(2021, 3, 5)]
now = datetime(2021, 3, 5, 22, 45, 0)


def test_keep_year():
    assert keep_year(now, dates, 2) == {datetime(2020, 4, 5), datetime(2021, 1, 2)}
    assert keep_year(now, dates, 0) == set()
    dates_gap = [datetime(2018, 2, 3), datetime(2019, 3, 4), datetime(2021, 5, 6)]
    assert keep_year(datetime(2021, 5, 6), dates_gap, 3) == {datetime(2019, 3, 4), datetime(2021, 5, 6)}
    assert keep_year(datetime(2022, 2, 3), dates_gap, 1) == set()


def test_keep_month():
    assert keep_month(now, dates, 3) == {datetime(2021, 1, 2), datetime(2021, 2, 2), datetime(2021, 3, 3)}


def test_keep_day():
    assert keep_day(now, dates, 3) == {datetime(2021, 3, 3), datetime(2021, 3, 4), datetime(2021, 3, 5)}


datetimes = [datetime(2021, 3, 5, 17, 0, 0), datetime(2021, 3, 5, 18, 0, 0), datetime(2021, 3, 5, 19, 0, 0),
             datetime(2021, 3, 5, 20, 0, 0), datetime(2021, 3, 5, 21, 0, 0), datetime(2021, 3, 5, 22, 0, 0),
             datetime(2021, 3, 5, 22, 15, 0), datetime(2021, 3, 5, 22, 30, 0), datetime(2021, 3, 5, 22, 45, 0)]


def test_keep_hour():
    assert keep_hour(now, datetimes, 3) == {datetime(2021, 3, 5, 20, 0, 0), datetime(2021, 3, 5, 21, 0, 0),
                                            datetime(2021, 3, 5, 22, 0, 0)}


def test_keep_frequent():
    assert keep_frequent(now, datetimes, 3) == {datetime(2021, 3, 5, 22, 15, 0), datetime(2021, 3, 5, 22, 30, 0),
                                                datetime(2021, 3, 5, 22, 45, 0)}