from datetime import datetime as dt
from retrosheet_data import date_math, weeks_between, date_from_int


def test_addition():
    start_date = "20240101"

    assert date_math(start_date, "+", year_delta=1) == "20250101"
    assert date_math(start_date, "+", month_delta=2) == "20240301"
    assert date_math(start_date, "+", day_delta=9) == "20240110"
    assert (
        date_math(start_date, "+", year_delta=1, month_delta=1, day_delta=1)
        == "20250202"
    )

    start_date = "20241201"

    assert date_math(start_date, "+", month_delta=1) == "20250101"
    assert date_math(start_date, "+", year_delta=2, month_delta=1) == "20270101"
    assert (
        date_math(start_date, "+", year_delta=2, month_delta=1, day_delta=10)
        == "20270111"
    )


def test_subtraction():
    start_date = "20241201"

    assert date_math(start_date, "-", year_delta=1) == "20231201"
    assert date_math(start_date, "-", month_delta=2) == "20241001"
    assert date_math(start_date, "-", day_delta=1) == "20241130"
    assert (
        date_math(start_date, "-", year_delta=1, month_delta=1, day_delta=1)
        == "20231031"
    )

    start_date = "20240101"
    assert date_math(start_date, "-", month_delta=2) == "20231101"
    assert date_math(start_date, "-", day_delta=2) == "20231230"


def test_date_from_int():
    assert date_from_int("20241201") == dt(2024, 12, 1)


def test_weeks_between():
    assert weeks_between("20241201", "20241214") == 2
    assert weeks_between("20241001", "20251001") == 52
    assert weeks_between("20251001", "20241001") == 52