session=session,
            )
            ti = dag_run.get_task_instance(dag_task.task_id, session)
            ti.set_state(TaskInstanceState.SUCCESS, session)
            session.flush()
            with mock.patch.object(Stats, "timing") as stats_mock:
                dag_run.update_state(session)
            metric_name = f"dagrun.{dag.dag_id}.first_task_scheduling_delay"
            if expected:
                true_delay = ti.start_date - dag_run.data_interval_end
                sched_delay_stat_call = call(metric_name, true_delay, tags=expected_stat_tags)
                sched_delay_stat_call_with_tags = call(
                    "dagrun.first_task_scheduling_delay", true_delay, tags=expected_stat_tags
                )
                assert (
                    sched_delay_stat_call in stats_mock.mock_calls
                    and sched_delay_stat_call_with_tags in stats_mock.mock_calls
                )
            else:
                
                sched_delay_stat_call = call(
                    metric_name,
                    mock.ANY,
                )
                assert sched_delay_stat_call not in stats_mock.mock_calls
        finally:
            
            session.rollback()
            session.close()
    def test_states_sets(self, session):