import pandas as pd

from ovretl.tracking_utils.extract_event_date import (
    extract_pickup_event_date,
    extract_arrival_event_date,
    add_tracking_event_dates,
)


def test_extract_pickup_event():
    pickup_events_shipment_df = pd.DataFrame.from_records(
        [
            {
                "shipment_id": 0,
                "event_description": "Pickup",
                "date_type": "actual",
                "date": pd.Timestamp("2020-01-02 23:00:00"),
                "location_type": "warehouse",
            },
            {
                "shipment_id": 0,
                "event_description": "Pickup",
                "date_type": "estimated",
                "date": pd.Timestamp("2015-01-05 23:00:00"),
                "location_type": "warehouse",
            },
        ]
    )
    pickup_date = extract_pickup_event_date(
        events_shipment_df=pickup_events_shipment_df,
    )
    assert pickup_date == pd.Timestamp("2015-01-05 23:00:00")


def test_extract_arrival_event():
    arrival_events_shipment_df = pd.DataFrame.from_records(
        [
            {
                "shipment_id": 0,
                "event_description": "Arrival",
                "date_type": "actual",
                "date": pd.Timestamp("2020-01-02 23:00:00"),
                "is_used_for_eta": False,
                "location_type": "harbor",
            },
            {
                "shipment_id": 0,
                "event_description": "Arrival",
                "date_type": "estimated",
                "date": pd.Timestamp("2015-01-05 23:00:00"),
                "is_used_for_eta": True,
                "location_type": "harbor",
            },
        ]
    )
    arrival_date = extract_arrival_event_date(
        events_shipment_df=arrival_events_shipment_df,
    )
    assert arrival_date == pd.Timestamp("2015-01-05 23:00:00")


def test_add_tracking_events():
    events_shipment_df = pd.DataFrame.from_records(
        [
            {
                "shipment_id": 0,
                "event_description": "Pickup",
                "date_type": "actual",
                "date": pd.Timestamp("2020-01-02 23:00:00"),
                "location_type": "warehouse",
                "is_used_for_etd": False,
                "is_used_for_eta": False,
            },
            {
                "shipment_id": 0,
                "event_description": "Pickup",
                "date_type": "estimated",
                "date": pd.Timestamp("2015-01-05 23:00:00"),
                "location_type": "warehouse",
                "is_used_for_etd": False,
                "is_used_for_eta": False,
            },
        ]
    )
    shipment = pd.Series(data=[0], index=["shipment_id"])
    assert add_tracking_event_dates(shipment, events_shipment_df).equals(shipment)
