Repository: apache/airflow
File: airflow-core/tests/unit/dag_processing/test_processor.py (Python)

Code before:
    )

    assert stat.last_finish_time == finish_time
    assert stat.run_count == 3
    assert stat.import_errors == 1


class TestExecuteCallbacks:
    def test_execute_callbacks_locks_bundle_version(self):
        request = DagCallbackRequest(
            filepath="test.py",
            dag_id="test_dag",
            run_id="test_run",
            bundle_name="testing",
            bundle_version="some_commit_hash",
            is_failure_callback=False,
            msg=None,
        )
        log = structlog.get_logger()
        dagbag = MagicMock()

        with (
            patch("airflow.dag_processing.processor.BundleVersionLock") as mock_lock,
            patch("airflow.dag_processing.processor._execute_dag_callbacks") as mock_execute,
        ):
            _execute_callbacks(dagbag, [request], log)

        mock_lock.assert_called_once_with(bundle_name="testing", bundle_version="some_commit_hash")
        mock_lock.return_value.__enter__.assert_called_once()
        mock_lock.return_value.__exit__.assert_called_once()
        mock_execute.assert_called_once_with(dagbag, request, log)


class TestExecuteDagCallbacks:
    """Test the _execute_dag_callbacks function with context_from_server"""

    def test_execute_dag_callbacks_with_context_from_server(self, spy_agency):
        """Test _execute_dag_callbacks uses RuntimeTaskInstance context when context_from_server is provided"""
        called = False
        context_received = None

        def on_failure(context):
            nonlocal called, context_received
            called = True
            context_received = context

        with DAG(dag_id="test_dag", on_failure_callback=on_failure) as dag:
            BaseOperator(task_id="test_task")

        def fake_collect_dags(self, *args, **kwargs):
            self.dags[dag.dag_id] = dag

Review comment (suggestion):
```suggestion
            _execute_callbacks(dagbag, callbacks, log)
```

Code after:
    )

    assert stat.last_finish_time == finish_time
    assert stat.run_count == 3
    assert stat.import_errors == 1


class TestExecuteCallbacks:
    def test_execute_callbacks_locks_bundle_version(self):
        callbacks = [
            DagCallbackRequest(
                filepath="test.py",
                dag_id="test_dag",
                run_id="test_run",
                bundle_name="testing",
                bundle_version="some_commit_hash",
                is_failure_callback=False,
                msg=None,
            )
        ]
        log = MagicMock(spec=FilteringBoundLogger)
        dagbag = MagicMock(spec=DagBag)

        with (
            patch("airflow.dag_processing.processor.BundleVersionLock") as mock_lock,
            patch("airflow.dag_processing.processor._execute_dag_callbacks") as mock_execute,
        ):
            _execute_callbacks(dagbag, callbacks, log)

        mock_lock.assert_called_once_with(bundle_name="testing", bundle_version="some_commit_hash")
        mock_lock.return_value.__enter__.assert_called_once()
        mock_lock.return_value.__exit__.assert_called_once()
        mock_execute.assert_called_once_with(dagbag, callbacks[0], log)


class TestExecuteDagCallbacks:
    """Test the _execute_dag_callbacks function with context_from_server"""

    def test_execute_dag_callbacks_with_context_from_server(self, spy_agency):
        """Test _execute_dag_callbacks uses RuntimeTaskInstance context when context_from_server is provided"""
        called = False
        context_received = None

        def on_failure(context):
            nonlocal called, context_received
            called = True
            context_received = context

        with DAG(dag_id="test_dag", on_failure_callback=on_failure) as dag:
            BaseOperator(task_id="test_task")