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

Code before:
                    bundle_path=dag2_path.bundle_path,
                    bundle_name="testing",
                    callbacks=[dag2_req1],
                    selector=mock.ANY,
                    logger=mock_logger,
                    logger_filehandle=mock_filehandle,
                    client=mock.ANY,
                ),
                mock.call(
                    id=mock.ANY,
                    path=Path(dag1_path.bundle_path, dag1_path.rel_path),
                    bundle_path=dag1_path.bundle_path,
                    bundle_name="testing",
                    callbacks=[dag1_req1, dag1_req2],
                    selector=mock.ANY,
                    logger=mock_logger,
                    logger_filehandle=mock_filehandle,
                    client=mock.ANY,
                ),
            ]
            # And removed from the queue
            assert dag1_path not in manager._callback_to_execute
            assert dag2_path not in manager._callback_to_execute

    @mock.patch("airflow.dag_processing.manager.DagBundlesManager")
    def test_add_callback_initializes_versioned_bundle(self, mock_bundle_manager):
        manager = DagFileProcessorManager(max_runs=1)
        bundle = MagicMock()
        bundle.supports_versioning = True
        bundle.path = Path("/tmp/bundle")
        mock_bundle_manager.return_value.get_bundle.return_value = bundle

        request = DagCallbackRequest(
            filepath="file1.py",
            dag_id="dag1",
            run_id="run1",
            is_failure_callback=False,
            bundle_name="testing",
            bundle_version="some_commit_hash",
            msg=None,
        )

        manager._add_callback_to_queue(request)

        bundle.initialize.assert_called_once()

    def test_dag_with_assets(self, session, configure_testing_dag_bundle):
        """'Integration' test to ensure that the assets get parsed and stored correctly for parsed dags."""
        test_dag_path = str(TEST_DAG_FOLDER / "test_assets.py")

        with configure_testing_dag_bundle(test_dag_path):

Review comment (bug):
This test verifies that `bundle.initialize()` is called for versioned bundles, but it doesn't cover the exception handling path where initialization fails. The implementation logs the exception and skips the callback. Add a test case where `bundle.initialize()` raises an exception to verify that the error is logged and the callback is not queued.

Code after:
                    bundle_path=dag2_path.bundle_path,
                    bundle_name="testing",
                    callbacks=[dag2_req1],
                    selector=mock.ANY,
                    logger=mock_logger,
                    logger_filehandle=mock_filehandle,
                    client=mock.ANY,
                ),
                mock.call(
                    id=mock.ANY,
                    path=Path(dag1_path.bundle_path, dag1_path.rel_path),
                    bundle_path=dag1_path.bundle_path,
                    bundle_name="testing",
                    callbacks=[dag1_req1, dag1_req2],
                    selector=mock.ANY,
                    logger=mock_logger,
                    logger_filehandle=mock_filehandle,
                    client=mock.ANY,
                ),
            ]
            # And removed from the queue
            assert dag1_path not in manager._callback_to_execute
            assert dag2_path not in manager._callback_to_execute

    @mock.patch("airflow.dag_processing.manager.DagBundlesManager")
    def test_add_callback_initializes_versioned_bundle(self, mock_bundle_manager):
        manager = DagFileProcessorManager(max_runs=1)
        bundle = MagicMock()
        bundle.supports_versioning = True
        bundle.path = Path("/tmp/bundle")
        mock_bundle_manager.return_value.get_bundle.return_value = bundle

        request = DagCallbackRequest(
            filepath="file1.py",
            dag_id="dag1",
            run_id="run1",
            is_failure_callback=False,
            bundle_name="testing",
            bundle_version="some_commit_hash",
            msg=None,
        )

        manager._add_callback_to_queue(request)

        bundle.initialize.assert_called_once()

    @mock.patch("airflow.dag_processing.manager.DagBundlesManager")
    def test_add_callback_skips_when_bundle_init_fails(self, mock_bundle_manager):
        manager = DagFileProcessorManager(max_runs=1)
        bundle = MagicMock()
        bundle.supports_versioning = True