"""
        Tests that adding State.failed_states and State.success_states work as expected.
        """
        dag = DAG(dag_id="test_dagrun_states", start_date=DEFAULT_DATE)
        dag_task_success = EmptyOperator(task_id="dummy", dag=dag)
        dag_task_failed = EmptyOperator(task_id="dummy2", dag=dag)
        initial_task_states = {
            dag_task_success.task_id: TaskInstanceState.SUCCESS,
            dag_task_failed.task_id: TaskInstanceState.FAILED,
        }
        dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session)
        ti_success = dag_run.get_task_instance(dag_task_success.task_id)
        ti_failed = dag_run.get_task_instance(dag_task_failed.task_id)
        assert ti_success.state in State.success_states
        assert ti_failed.state in State.failed_states
@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, session, run_type, expected_tis):