diff --git a/airflow/models/dagbag.py b/airflow/models/dagbag.py index 8a5557c2c2264..cabc142f31708 100644 --- a/airflow/models/dagbag.py +++ b/airflow/models/dagbag.py @@ -54,6 +54,7 @@ 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 @@ def __init__( 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 0a2c062c66436..67de520c38b5d 100644 --- a/tests/cli/conftest.py +++ b/tests/cli/conftest.py @@ -24,6 +24,7 @@ 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 @@ 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 4284bf02d9695..81fd6a3017797 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 @@ def test_find_zombies_nothing(self): 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 @@ def test_find_zombies(self): 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 """