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

Code before:
        dag_run.update_state(session=session)

        mock_get_by_id.assert_called_once_with(deadline_id, session)
        mock_prune.assert_not_called()
        assert dag_run.state == DagRunState.SUCCESS

    @mock.patch.object(Deadline, "prune_deadlines")
    def test_dagrun_success_handles_empty_deadline_list(self, mock_prune, dag_maker, session):
        with dag_maker(
            dag_id="test_dagrun_empty_deadlines",
            schedule=datetime.timedelta(days=1),
        ) as dag:
            EmptyOperator(task_id="task_1")

        initial_task_states = {"task_1": TaskInstanceState.SUCCESS}
        dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session)
        dag_run = session.merge(dag_run)
        dag.deadline = []
        dag_run.dag = dag

        dag_run.update_state(session=session)

        mock_prune.assert_not_called()
        assert dag_run.state == DagRunState.SUCCESS

    def test_clear_task_instances_recalculates_dagrun_queued_deadlines(self, session, dag_maker):
        """Test that clearing tasks recalculates all (and only) DAGRUN_QUEUED_AT deadlines."""
        with dag_maker(
            dag_id="test_recalculate_deadlines",
            schedule=datetime.timedelta(days=1),
        ) as dag:
            EmptyOperator(task_id="task_1")

        dag_run = self.create_dag_run(
            dag=dag,
            task_states={"task_1": TaskInstanceState.SUCCESS},
            session=session,
        )

        original_queued_at = timezone.utcnow() - datetime.timedelta(hours=2)
        dag_run.queued_at = original_queued_at
        session.flush()

        serialized_dag_id = session.scalar(
            select(SerializedDagModel.id).where(SerializedDagModel.dag_id == dag.dag_id)
        )

        deadline_configs = [
            (DeadlineReference.DAGRUN_QUEUED_AT, datetime.timedelta(hours=1)),
            (DeadlineReference.DAGRUN_QUEUED_AT, datetime.timedelta(hours=2)),
            (DeadlineReference.FIXED_DATETIME, datetime.timedelta(hours=1)),

Review comment (suggestion):
Should not it be in `test_taskinstance`? You updated `taskinstance.py` so I expect the test to be in `test_taskinstance`

Code after:

        mock_get_by_id.assert_called_once_with(deadline_id, session)
        mock_prune.assert_not_called()
        assert dag_run.state == DagRunState.SUCCESS

    @mock.patch.object(Deadline, "prune_deadlines")
    def test_dagrun_success_handles_empty_deadline_list(self, mock_prune, dag_maker, session):
        with dag_maker(
            dag_id="test_dagrun_empty_deadlines",
            schedule=datetime.timedelta(days=1),
        ) as dag:
            EmptyOperator(task_id="task_1")

        initial_task_states = {"task_1": TaskInstanceState.SUCCESS}
        dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session)
        dag_run = session.merge(dag_run)
        dag.deadline = []
        dag_run.dag = dag

        dag_run.update_state(session=session)

        mock_prune.assert_not_called()
        assert dag_run.state == DagRunState.SUCCESS


@pytest.mark.parametrize(
    ("run_type", "expected_tis"),
    [
        pytest.param(DagRunType.MANUAL, 1, id="manual"),
        pytest.param(DagRunType.BACKFILL_JOB, 3, id="backfill"),
    ],
)
@mock.patch.object(Stats, "incr")
def test_verify_integrity_task_start_and_end_date(Stats_incr, dag_maker, session, run_type, expected_tis):
    """Test that tasks with specific dates are only created for backfill runs"""
    with dag_maker("test", schedule=datetime.timedelta(days=1), start_date=DEFAULT_DATE) as dag:
        EmptyOperator(task_id="without")
        EmptyOperator(task_id="with_start_date", start_date=DEFAULT_DATE + datetime.timedelta(1))
        EmptyOperator(task_id="with_end_date", end_date=DEFAULT_DATE - datetime.timedelta(1))

    dag_run = DagRun(
        dag_id=dag.dag_id,
        run_type=run_type,
        logical_date=DEFAULT_DATE,
        run_id=DagRun.generate_run_id(run_type=run_type, logical_date=DEFAULT_DATE, run_after=DEFAULT_DATE),
    )
    dag_run.dag = dag

    session.add(dag_run)
    session.flush()
    dag_version_id = DagVersion.get_latest_version(dag.dag_id, session=session).id