from datetime import datetime, timedelta

from airflow.operators.dummy_operator import DummyOperator

from helpers.dag import DAG
from helpers.lineage.sensors import LineageSensor

dag = (
    DAG()
    .airflow()
    .defaults(
        start_date=datetime(2019, 5, 27),
        retries=0,
        retry_delay=timedelta(minutes=1),
    )
    .entities(
        {
            "entity_test": {
                "sensor": "task_sensor_test",
                "generator": "task_generate_test",
                "generated_states": ["completed"]
                },
        }
    )
    .schedule("0 0 * * *")
    .build()
)

with dag:

    sensor_test = LineageSensor(task_id="task_sensor_test")

    entity_test = DummyOperator(task_id="task_generate_test")

    entity_test.set_upstream(sensor_test)