Repository: apache/airflow
File: providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py (Python)

Code before:
        self.dag_model = DagModel(dag_id=TRIGGERED_DAG_ID, fileloc=self._tmpfile)
        with create_session() as session:
            session.add(self.dag_model)
            session.commit()

    def teardown_method(self):
        """Cleanup state after testing in DB."""
        with create_session() as session:
            session.execute(delete(Log).where(Log.dag_id == TEST_DAG_ID))
            for dbmodel in [DagModel, DagRun, TaskInstance]:
                session.execute(delete(dbmodel).where(dbmodel.dag_id.in_([TRIGGERED_DAG_ID, TEST_DAG_ID])))

    def test_trigger_dagrun(self, dag_maker, mock_supervisor_comms):
        """Test TriggerDagRunOperator."""
        with time_machine.travel("2025-02-18T08:04:46Z", tick=False):
            with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
                task = TriggerDagRunOperator(
                    task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, note="Test note"
                )
            mock_warning = mock.patch.object(task.log, "warning").start()
            dag_maker.sync_dagbag_to_db()
            parse_and_sync_to_db(self.f_name)
            dag_maker.create_dagrun()
            task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)

            mock_warning.assert_called_once_with(
                "Parameter 'note' is not supported in Airflow 2.x and will be ignored."
            )
            dagrun = dag_maker.session.scalar(select(DagRun).where(DagRun.dag_id == TRIGGERED_DAG_ID))
            assert dagrun.run_type == DagRunType.MANUAL
            assert dagrun.run_id == DagRun.generate_run_id(DagRunType.MANUAL, dagrun.logical_date)

    def test_explicitly_provided_trigger_run_id_is_saved_as_attr(self, dag_maker, session):
        with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
            task = TriggerDagRunOperator(
                task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, trigger_run_id="test_run_id"
            )
            assert task.trigger_run_id == "test_run_id"
        dag_maker.create_dagrun()
        task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)

        assert task.trigger_run_id == "test_run_id"

    def test_extra_operator_link(self, dag_maker, session):
        """Asserts whether the correct extra links url will be created."""
        with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
            task = TriggerDagRunOperator(
                task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, trigger_run_id="test_run_id"
            )
        dag_maker.create_dagrun()
        task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)

Review comment (suggestion):
Can you please restructure this assertion to use the assert keyword? This should look like:

```python
def test_trigger_dagrun(self, dag_maker, mock_supervisor_comms, mocker):
...
        # If we want to assert that one of multiple calls was the right one:
        assert mocker.call(...) in mock_warning.mock_calls

        # If we want to assert there was exactly one call, and it was the right one:
        assert mock_warning.mock_calls == [mocker.call(...)]

        # If we want to assert there were no calls:
        assert mock_warning.mock_calls == []

```

See also: https://lists.apache.org/thread/41b04mg0rolv0sj98jhogsztstxnqfg5

Code after:
        self.dag_model = DagModel(dag_id=TRIGGERED_DAG_ID, fileloc=self._tmpfile)
        with create_session() as session:
            session.add(self.dag_model)
            session.commit()

    def teardown_method(self):
        """Cleanup state after testing in DB."""
        with create_session() as session:
            session.execute(delete(Log).where(Log.dag_id == TEST_DAG_ID))
            for dbmodel in [DagModel, DagRun, TaskInstance]:
                session.execute(delete(dbmodel).where(dbmodel.dag_id.in_([TRIGGERED_DAG_ID, TEST_DAG_ID])))

    def test_trigger_dagrun(self, dag_maker, mock_supervisor_comms):
        """Test TriggerDagRunOperator."""
        with time_machine.travel("2025-02-18T08:04:46Z", tick=False):
            with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
                task = TriggerDagRunOperator(
                    task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, note="Test note"
                )
            mock_warning = mock.patch.object(task.log, "warning").start()
            dag_maker.sync_dagbag_to_db()
            parse_and_sync_to_db(self.f_name)
            dag_maker.create_dagrun()
            task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)

            dagrun = dag_maker.session.scalar(select(DagRun).where(DagRun.dag_id == TRIGGERED_DAG_ID))
            assert mock_warning.mock_calls == [
                mock.call("Parameter 'note' is not supported in Airflow 2.x and will be ignored.")
            ]
            assert dagrun.run_type == DagRunType.MANUAL
            assert dagrun.run_id == DagRun.generate_run_id(DagRunType.MANUAL, dagrun.logical_date)

    def test_explicitly_provided_trigger_run_id_is_saved_as_attr(self, dag_maker, session):
        with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
            task = TriggerDagRunOperator(
                task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, trigger_run_id="test_run_id"
            )
            assert task.trigger_run_id == "test_run_id"
        dag_maker.create_dagrun()
        task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)

        assert task.trigger_run_id == "test_run_id"

    def test_extra_operator_link(self, dag_maker, session):
        """Asserts whether the correct extra links url will be created."""
        with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True):
            task = TriggerDagRunOperator(
                task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, trigger_run_id="test_run_id"
            )
        dag_maker.create_dagrun()
        task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)