Remove side-effect of Hybrid executor blocking test (#39615)

The tests introduced in #39531 introduced side effect of cleaning the dictionary of executors - they started to fail main in tests where bothn DB and NonDB tests were run together.
apache · May 14, 2024 · af3e2cf · af3e2cf
1 parent 671d1e8
commit af3e2cf
Showing 1 changed file with 28 additions and 14 deletions.
diff --git a/tests/executors/test_executor_loader.py b/tests/executors/test_executor_loader.py
@@ -19,6 +19,7 @@
 from contextlib import nullcontext
 from importlib import reload
 from unittest import mock
+from unittest.mock import patch
 
 import pytest
 
@@ -55,6 +56,12 @@ def setup_method(self) -> None:
         global ExecutorLoader
         ExecutorLoader = executor_loader.ExecutorLoader  # type: ignore
 
+    def teardown_method(self) -> None:
+        from airflow.executors import executor_loader
+
+        reload(executor_loader)
+        ExecutorLoader.init_executors()
+
     def test_no_executor_configured(self):
         with conf_vars({("core", "executor"): None}):
             with pytest.raises(AirflowConfigException, match=r".*not found in config$"):
@@ -304,19 +311,26 @@ def test_validate_database_executor_compatibility_sqlite(self, monkeypatch, exec
             ExecutorLoader.validate_database_executor_compatibility(executor)
 
     def test_load_executor(self):
-        ExecutorLoader.block_use_of_hybrid_exec = mock.Mock()
-        with conf_vars({("core", "executor"): "LocalExecutor"}):
-            ExecutorLoader.init_executors()
-            assert isinstance(ExecutorLoader.load_executor("LocalExecutor"), LocalExecutor)
-            assert isinstance(ExecutorLoader.load_executor(executor_loader._executor_names[0]), LocalExecutor)
-            assert isinstance(ExecutorLoader.load_executor(None), LocalExecutor)
+        with patch.object(ExecutorLoader, "block_use_of_hybrid_exec"):
+            with conf_vars({("core", "executor"): "LocalExecutor"}):
+                ExecutorLoader.init_executors()
+                assert isinstance(ExecutorLoader.load_executor("LocalExecutor"), LocalExecutor)
+                assert isinstance(
+                    ExecutorLoader.load_executor(executor_loader._executor_names[0]), LocalExecutor
+                )
+                assert isinstance(ExecutorLoader.load_executor(None), LocalExecutor)
 
     def test_load_executor_alias(self):
-        ExecutorLoader.block_use_of_hybrid_exec = mock.Mock()
-        with conf_vars({("core", "executor"): "local_exec:airflow.executors.local_executor.LocalExecutor"}):
-            ExecutorLoader.init_executors()
-            assert isinstance(ExecutorLoader.load_executor("local_exec"), LocalExecutor)
-            assert isinstance(
-                ExecutorLoader.load_executor("airflow.executors.local_executor.LocalExecutor"), LocalExecutor
-            )
-            assert isinstance(ExecutorLoader.load_executor(executor_loader._executor_names[0]), LocalExecutor)
+        with patch.object(ExecutorLoader, "block_use_of_hybrid_exec"):
+            with conf_vars(
+                {("core", "executor"): "local_exec:airflow.executors.local_executor.LocalExecutor"}
+            ):
+                ExecutorLoader.init_executors()
+                assert isinstance(ExecutorLoader.load_executor("local_exec"), LocalExecutor)
+                assert isinstance(
+                    ExecutorLoader.load_executor("airflow.executors.local_executor.LocalExecutor"),
+                    LocalExecutor,
+                )
+                assert isinstance(
+                    ExecutorLoader.load_executor(executor_loader._executor_names[0]), LocalExecutor
+                )