Ensure dag.test uses serialized dag for testing

airflow-core/tests/unit/models/test_dag.py

@@ -53,6 +53,7 @@
     get_asset_triggered_next_run_info,
     get_next_data_interval,
 )
+from airflow.models.dagbag import DBDagBag
 from airflow.models.dagbundle import DagBundleModel
 from airflow.models.dagrun import DagRun
 from airflow.models.serialized_dag import SerializedDagModel
@@ -80,6 +81,7 @@
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
 
 from tests_common.test_utils.asserts import assert_queries_count
+from tests_common.test_utils.config import conf_vars
 from tests_common.test_utils.dag import create_scheduler_dag, sync_dag_to_db
 from tests_common.test_utils.db import (
     clear_db_assets,
@@ -179,6 +181,29 @@ def setup_method(self) -> None:
         clear_db_dags()
         clear_db_assets()
 
+    @conf_vars({("core", "load_examples"): "false"})
+    def test_dag_test_auto_parses_when_not_serialized(self, test_dags_bundle, session):
+        """
+        DAG.test() should auto-parse and sync the DAG if it's not serialized yet.
+        """
+
+        dag_id = "test_example_bash_operator"
+
+        dagbag = DagBag(dag_folder=os.fspath(TEST_DAGS_FOLDER), include_examples=False)
+        dag = dagbag.dags.get(dag_id)
+
+        # Ensure not serialized yet
+        assert DBDagBag().get_latest_version_of_dag(dag_id, session=session) is None
+        assert session.scalar(select(DagRun).where(DagRun.dag_id == dag_id)) is None
+
+        dr = dag.test()
+        assert dr is not None
+
+        # Serialized DAG should now exist and DagRun would be created
+        ser = DBDagBag().get_latest_version_of_dag(dag_id, session=session)
+        assert ser is not None
+        assert session.scalar(select(DagRun).where(DagRun.dag_id == dag_id)) is not None
+
     def teardown_method(self) -> None:
         clear_db_runs()
         clear_db_dags()

task-sdk/src/airflow/sdk/definitions/dag.py

@@ -1197,7 +1197,33 @@ def test(
             data_interval = (
                 self.timetable.infer_manual_data_interval(run_after=logical_date) if logical_date else None
             )
-            scheduler_dag = SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(self))  # type: ignore[arg-type]
+            from airflow.models.dag_version import DagVersion
+
+            version = DagVersion.get_version(self.dag_id)
+            if not version:
+                from airflow.dag_processing.bundles.manager import DagBundlesManager
+                from airflow.dag_processing.dagbag import DagBag, sync_bag_to_db
+                from airflow.sdk.definitions._internal.dag_parsing_context import (
+                    _airflow_parsing_context_manager,
+                )
+
+                manager = DagBundlesManager()
+                manager.sync_bundles_to_db(session=session)
+                session.commit()
+                # sync all bundles? or use the dags-folder bundle?
+                # What if the test dag is in a different bundle?
+                for bundle in manager.get_all_dag_bundles():
+                    if not bundle.is_initialized:
+                        bundle.initialize()
+                    with _airflow_parsing_context_manager(dag_id=self.dag_id):
+                        dagbag = DagBag(
+                            dag_folder=bundle.path, bundle_path=bundle.path, include_examples=False
+                        )
+                        sync_bag_to_db(dagbag, bundle.name, bundle.version)
+                    version = DagVersion.get_version(self.dag_id)
+                    if version:
+                        break
+            scheduler_dag = SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(self))
             # Preserve callback functions from original Dag since they're lost during serialization
             # and yes it is a hack for now! It is a tradeoff for code simplicity.
             # Without it, we need "Scheduler Dag" (Serialized dag) for the scheduler bits
@@ -1206,8 +1232,8 @@ def test(
 
             # Scheduler DAG shouldn't have these attributes, but assigning them
             # here is an easy hack to get this test() thing working.
-            scheduler_dag.on_success_callback = self.on_success_callback  # type: ignore[attr-defined]
-            scheduler_dag.on_failure_callback = self.on_failure_callback  # type: ignore[attr-defined]
+            scheduler_dag.on_success_callback = self.on_success_callback  # type: ignore[attr-defined, union-attr]
+            scheduler_dag.on_failure_callback = self.on_failure_callback  # type: ignore[attr-defined, union-attr]
 
             dr: DagRun = get_or_create_dagrun(
                 dag=scheduler_dag,