From 3f1220d5b998ec691ac081d0f58a19bc2963db97 Mon Sep 17 00:00:00 2001
From: Amogh Desai <amoghrajesh1999@gmail.com>
Date: Fri, 23 Aug 2024 11:17:06 +0530
Subject: [PATCH] Splitting syspath preparation into stages (#41672)

(cherry picked from commit 0be5decdc8284f4f2f3d6be66277cc102c842b5c)
---
 airflow/settings.py         | 17 ++++++++++-------
 tests/core/test_settings.py | 33 +++++++++++++++++++++++++++------
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/airflow/settings.py b/airflow/settings.py
index 175a63f69d2f4..dc24a2c5acc5a 100644
--- a/airflow/settings.py
+++ b/airflow/settings.py
@@ -675,11 +675,8 @@ def configure_action_logging() -> None:
     """Any additional configuration (register callback) for airflow.utils.action_loggers module."""
 
 
-def prepare_syspath():
-    """Ensure certain subfolders of AIRFLOW_HOME are on the classpath."""
-    if DAGS_FOLDER not in sys.path:
-        sys.path.append(DAGS_FOLDER)
-
+def prepare_syspath_for_config_and_plugins():
+    """Update sys.path for the config and plugins directories."""
     # Add ./config/ for loading custom log parsers etc, or
     # airflow_local_settings etc.
     config_path = os.path.join(AIRFLOW_HOME, "config")
@@ -690,6 +687,12 @@ def prepare_syspath():
         sys.path.append(PLUGINS_FOLDER)
 
 
+def prepare_syspath_for_dags_folder():
+    """Update sys.path to include the DAGs folder."""
+    if DAGS_FOLDER not in sys.path:
+        sys.path.append(DAGS_FOLDER)
+
+
 def get_session_lifetime_config():
     """Get session timeout configs and handle outdated configs gracefully."""
     session_lifetime_minutes = conf.get("webserver", "session_lifetime_minutes", fallback=None)
@@ -771,12 +774,13 @@ def import_local_settings():
 def initialize():
     """Initialize Airflow with all the settings from this file."""
     configure_vars()
-    prepare_syspath()
+    prepare_syspath_for_config_and_plugins()
     configure_policy_plugin_manager()
     # Load policy plugins _before_ importing airflow_local_settings, as Pluggy uses LIFO and we want anything
     # in airflow_local_settings to take precendec
     load_policy_plugins(POLICY_PLUGIN_MANAGER)
     import_local_settings()
+    prepare_syspath_for_dags_folder()
     global LOGGING_CLASS_PATH
     LOGGING_CLASS_PATH = configure_logging()
     State.state_color.update(STATE_COLORS)
@@ -806,7 +810,6 @@ def is_usage_data_collection_enabled() -> bool:
 MEGABYTE = KILOBYTE * KILOBYTE
 WEB_COLORS = {"LIGHTBLUE": "#4d9de0", "LIGHTORANGE": "#FF9933"}
 
-
 # Updating serialized DAG can not be faster than a minimum interval to reduce database
 # write rate.
 MIN_SERIALIZED_DAG_UPDATE_INTERVAL = conf.getint("core", "min_serialized_dag_update_interval", fallback=30)
diff --git a/tests/core/test_settings.py b/tests/core/test_settings.py
index d05344bfa91d8..483ef24e25f7f 100644
--- a/tests/core/test_settings.py
+++ b/tests/core/test_settings.py
@@ -115,21 +115,42 @@ def teardown_method(self):
         for mod in [m for m in sys.modules if m not in self.old_modules]:
             del sys.modules[mod]
 
+    @mock.patch("airflow.settings.prepare_syspath_for_config_and_plugins")
     @mock.patch("airflow.settings.import_local_settings")
-    @mock.patch("airflow.settings.prepare_syspath")
-    def test_initialize_order(self, prepare_syspath, import_local_settings):
+    @mock.patch("airflow.settings.prepare_syspath_for_dags_folder")
+    def test_initialize_order(
+        self,
+        mock_prepare_syspath_for_dags_folder,
+        mock_import_local_settings,
+        mock_prepare_syspath_for_config_and_plugins,
+    ):
         """
-        Tests that import_local_settings is called after prepare_classpath
+        Tests that import_local_settings is called between prepare_syspath_for_config_and_plugins
+        and prepare_syspath_for_dags_folder
         """
         mock_local_settings = mock.Mock()
-        mock_local_settings.attach_mock(prepare_syspath, "prepare_syspath")
-        mock_local_settings.attach_mock(import_local_settings, "import_local_settings")
+
+        mock_local_settings.attach_mock(
+            mock_prepare_syspath_for_config_and_plugins, "prepare_syspath_for_config_and_plugins"
+        )
+        mock_local_settings.attach_mock(mock_import_local_settings, "import_local_settings")
+        mock_local_settings.attach_mock(
+            mock_prepare_syspath_for_dags_folder, "prepare_syspath_for_dags_folder"
+        )
 
         import airflow.settings
 
         airflow.settings.initialize()
 
-        mock_local_settings.assert_has_calls([call.prepare_syspath(), call.import_local_settings()])
+        expected_calls = [
+            call.prepare_syspath_for_config_and_plugins(),
+            call.import_local_settings(),
+            call.prepare_syspath_for_dags_folder(),
+        ]
+
+        mock_local_settings.assert_has_calls(expected_calls)
+
+        assert mock_local_settings.mock_calls == expected_calls
 
     def test_import_with_dunder_all_not_specified(self):
         """