(airflow) branch main updated: Move tests `bootstrap_dagbag` into test utils (#44371)

[rom]

This is an automated email from the ASF dual-hosted git repository.

rom pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new f0af1b3d23b Move tests `bootstrap_dagbag` into test utils (#44371)
f0af1b3d23b is described below

commit f0af1b3d23bd9a0adfe9642cae8a2725435e1ccf
Author: Jed Cunningham <66968678+jedcunning...@users.noreply.github.com>
AuthorDate: Tue Nov 26 01:01:41 2024 -0700

    Move tests `bootstrap_dagbag` into test utils (#44371)
---
 airflow/utils/db.py           | 13 -------------
 tests_common/test_utils/db.py | 15 ++++++++++++++-
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 1c17e1d00f5..6c69bf2abcb 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -1224,19 +1224,6 @@ def resetdb(session: Session = NEW_SESSION, skip_init: 
bool = False):
         initdb(session=session)
 
 
-@provide_session
-def bootstrap_dagbag(session: Session = NEW_SESSION):
-    from airflow.models.dag import DAG
-    from airflow.models.dagbag import DagBag
-
-    dagbag = DagBag()
-    # Save DAGs in the ORM
-    dagbag.sync_to_db(session=session)
-
-    # Deactivate the unknown ones
-    DAG.deactivate_unknown_dags(dagbag.dags.keys(), session=session)
-
-
 @provide_session
 def downgrade(*, to_revision, from_revision=None, show_sql_only=False, 
session: Session = NEW_SESSION):
     """
diff --git a/tests_common/test_utils/db.py b/tests_common/test_utils/db.py
index c9b62d096b9..bf4afd260f6 100644
--- a/tests_common/test_utils/db.py
+++ b/tests_common/test_utils/db.py
@@ -54,6 +54,19 @@ from tests_common.test_utils.compat import (
 )
 
 
+def _bootstrap_dagbag():
+    from airflow.models.dag import DAG
+    from airflow.models.dagbag import DagBag
+
+    with create_session() as session:
+        dagbag = DagBag()
+        # Save DAGs in the ORM
+        dagbag.sync_to_db(session=session)
+
+        # Deactivate the unknown ones
+        DAG.deactivate_unknown_dags(dagbag.dags.keys(), session=session)
+
+
 def initial_db_init():
     from flask import Flask
 
@@ -63,7 +76,7 @@ def initial_db_init():
     from airflow.www.extensions.init_auth_manager import get_auth_manager
 
     db.resetdb()
-    db.bootstrap_dagbag()
+    _bootstrap_dagbag()
     # minimal app to add roles
     flask_app = Flask(__name__)
     flask_app.config["SQLALCHEMY_DATABASE_URI"] = conf.get("database", 
"SQL_ALCHEMY_CONN")