[airflow] branch main updated: Add fixture for CLI tests requiring sample dags (#26536)

[ds...@apache.org]

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

dstandish 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 3cd4df16d4 Add fixture for CLI tests requiring sample dags (#26536)
3cd4df16d4 is described below

commit 3cd4df16d4f383c27f7fc6bd932bca1f83ab9977
Author: Daniel Standish <15...@users.noreply.github.com>
AuthorDate: Tue Sep 20 21:37:18 2022 -0700

    Add fixture for CLI tests requiring sample dags (#26536)
    
    I noticed that various CLI tests were failing (specifically in test_task_command.py) when run locally but not when run in breeze.  I figured out the cause is that breeze has ('core', 'load_examples') set to true but I had it false locally.  To fix this we can add a fixture that patches the conf settings.  But we have to go one step further.  If we grab the conf value in the default for the argument in DagBag, this means the value is fixed at the time the class gets defined.  So the va [...]
    
    We do the same for safe_mode while we're at it, because why not.
---
 airflow/models/dagbag.py         | 14 ++++++++++++--
 tests/cli/conftest.py            |  7 +++++++
 tests/jobs/test_scheduler_job.py | 10 ++++++++--
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/airflow/models/dagbag.py b/airflow/models/dagbag.py
index 8a5557c2c2..cabc142f31 100644
--- a/airflow/models/dagbag.py
+++ b/airflow/models/dagbag.py
@@ -54,6 +54,7 @@ from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.retries import MAX_DB_RETRIES, run_with_db_retries
 from airflow.utils.session import provide_session
 from airflow.utils.timeout import timeout
+from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     import pathlib
@@ -92,8 +93,8 @@ class DagBag(LoggingMixin):
     def __init__(
         self,
         dag_folder: str | pathlib.Path | None = None,
-        include_examples: bool = conf.getboolean('core', 'LOAD_EXAMPLES'),
-        safe_mode: bool = conf.getboolean('core', 'DAG_DISCOVERY_SAFE_MODE'),
+        include_examples: bool | ArgNotSet = NOTSET,
+        safe_mode: bool | ArgNotSet = NOTSET,
         read_dags_from_db: bool = False,
         store_serialized_dags: bool | None = None,
         load_op_links: bool = True,
@@ -103,6 +104,15 @@ class DagBag(LoggingMixin):
 
         super().__init__()
 
+        include_examples = (
+            include_examples
+            if isinstance(include_examples, bool)
+            else conf.getboolean('core', 'LOAD_EXAMPLES')
+        )
+        safe_mode = (
+            safe_mode if isinstance(safe_mode, bool) else conf.getboolean('core', 'DAG_DISCOVERY_SAFE_MODE')
+        )
+
         if store_serialized_dags:
             warnings.warn(
                 "The store_serialized_dags parameter has been deprecated. "
diff --git a/tests/cli/conftest.py b/tests/cli/conftest.py
index 0a2c062c66..67de520c38 100644
--- a/tests/cli/conftest.py
+++ b/tests/cli/conftest.py
@@ -24,6 +24,7 @@ import pytest
 from airflow import models
 from airflow.cli import cli_parser
 from airflow.executors import celery_executor, celery_kubernetes_executor
+from tests.test_utils.config import conf_vars
 
 # Create custom executors here because conftest is imported first
 custom_executor_module = type(sys)('custom_executor')
@@ -36,6 +37,12 @@ custom_executor_module.CustomCeleryKubernetesExecutor = type(  # type: ignore
 sys.modules['custom_executor'] = custom_executor_module
 
 
+@pytest.fixture(autouse=True)
+def load_examples():
+    with conf_vars({('core', 'load_examples'): 'True'}):
+        yield
+
+
 @pytest.fixture(scope="session")
 def dagbag():
     return models.DagBag(include_examples=True)
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 4284bf02d9..81fd6a3017 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -104,6 +104,12 @@ def dagbag():
     return DagBag(read_dags_from_db=True)
 
 
+@pytest.fixture
+def load_examples():
+    with conf_vars({('core', 'load_examples'): 'True'}):
+        yield
+
+
 @pytest.mark.usefixtures("disable_load_example")
 @pytest.mark.need_serialized_dag
 class TestSchedulerJob:
@@ -4014,7 +4020,7 @@ class TestSchedulerJob:
 
             self.scheduler_job.executor.callback_sink.send.assert_not_called()
 
-    def test_find_zombies(self):
+    def test_find_zombies(self, load_examples):
         dagbag = DagBag(TEST_DAG_FOLDER, read_dags_from_db=False)
         with create_session() as session:
             session.query(LocalTaskJob).delete()
@@ -4072,7 +4078,7 @@ class TestSchedulerJob:
             session.query(TaskInstance).delete()
             session.query(LocalTaskJob).delete()
 
-    def test_zombie_message(self):
+    def test_zombie_message(self, load_examples):
         """
         Check that the zombie message comes out as expected
         """