Repository: apache/airflow
File: airflow-core/tests/unit/models/test_dag.py (Python)

Code before:
            deadline=DeadlineAlert(
                reference=reference_type,
                interval=interval,
                callback=AsyncCallback(empty_callback_for_deadline),
            ),
        )

        # Sync the DAG to the database to create DeadlineAlert records
        scheduler_dag = sync_dag_to_db(dag, session=session)

        dr = scheduler_dag.create_dagrun(
            run_id="test_dagrun_deadline",
            run_type=DagRunType.SCHEDULED,
            state=State.QUEUED,
            logical_date=TEST_DATE,
            run_after=TEST_DATE,
            triggered_by=DagRunTriggeredByType.TEST,
        )
        session.flush()
        dr = session.merge(dr)

        assert len(dr.deadlines) == 1
        assert dr.deadlines[0].deadline_time == getattr(dr, reference_column, DEFAULT_DATE) + interval

    def test_dag_with_multiple_deadlines(self, testing_dag_bundle, session):
        """Test that a DAG with multiple deadlines stores all deadlines and persists on re-serialization."""
        deadlines = [
            DeadlineAlert(
                reference=DeadlineReference.DAGRUN_QUEUED_AT,
                interval=datetime.timedelta(minutes=5),
                callback=AsyncCallback(empty_callback_for_deadline),
            ),
            DeadlineAlert(
                reference=DeadlineReference.DAGRUN_QUEUED_AT,
                interval=datetime.timedelta(minutes=10),
                callback=AsyncCallback(empty_callback_for_deadline),
            ),
            DeadlineAlert(
                reference=DeadlineReference.DAGRUN_LOGICAL_DATE,
                interval=datetime.timedelta(hours=1),
                callback=AsyncCallback(empty_callback_for_deadline),
            ),
        ]
        expected_num_deadlines = 3

        dag = DAG(
            dag_id="test_multiple_deadlines",
            schedule=datetime.timedelta(days=1),
            deadline=deadlines,
        )


Review comment (suggestion):
```suggestion
        """Test that a Dag with multiple deadlines stores all deadlines and persists on re-serialization."""
```

Code after:
        assert not orm_dag_owners

    @pytest.mark.need_serialized_dag
    @pytest.mark.parametrize(
        ("reference_type", "reference_column"),
        [
            pytest.param(DeadlineReference.DAGRUN_LOGICAL_DATE, "logical_date", id="logical_date"),
            pytest.param(DeadlineReference.DAGRUN_QUEUED_AT, "queued_at", id="queued_at"),
            pytest.param(DeadlineReference.FIXED_DATETIME(DEFAULT_DATE), "NONE", id="fixed_deadline"),
        ],
    )
    def test_dagrun_deadline(self, reference_type, reference_column, testing_dag_bundle, session):
        interval = datetime.timedelta(hours=1)
        dag = DAG(
            dag_id="test_deadline",
            schedule=datetime.timedelta(days=1),
            deadline=DeadlineAlert(
                reference=reference_type,
                interval=interval,
                callback=AsyncCallback(empty_callback_for_deadline),
            ),
        )

        # Sync the DAG to the database to create DeadlineAlert records
        scheduler_dag = sync_dag_to_db(dag, session=session)

        dr = scheduler_dag.create_dagrun(
            run_id="test_dagrun_deadline",
            run_type=DagRunType.SCHEDULED,
            state=State.QUEUED,
            logical_date=TEST_DATE,
            run_after=TEST_DATE,
            triggered_by=DagRunTriggeredByType.TEST,
        )
        session.flush()
        dr = session.merge(dr)

        assert len(dr.deadlines) == 1
        assert dr.deadlines[0].deadline_time == getattr(dr, reference_column, DEFAULT_DATE) + interval

    def test_dag_with_multiple_deadlines(self, testing_dag_bundle, session):
        """Test that a Dag with multiple deadlines stores all deadlines and persists on re-serialization."""
        deadlines = [
            DeadlineAlert(
                reference=DeadlineReference.DAGRUN_QUEUED_AT,
                interval=datetime.timedelta(minutes=5),
                callback=AsyncCallback(empty_callback_for_deadline),
            ),
            DeadlineAlert(
                reference=DeadlineReference.DAGRUN_QUEUED_AT,
                interval=datetime.timedelta(minutes=10),