From d25f00a57320834004631c03d6f8f25b604fe205 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Tue, 15 Aug 2023 15:26:28 -0700 Subject: [PATCH 1/9] add a mechanism to warn if executors override existing CLI commands --- airflow/cli/cli_parser.py | 16 ++++++++++++++++ tests/cli/test_cli_parser.py | 21 ++++++++++++++++++++- 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index 8e4d819098c5a..489a97f59e616 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -65,6 +65,22 @@ ALL_COMMANDS_DICT: dict[str, CLICommand] = {sp.name: sp for sp in airflow_commands} +# Check if sub-commands are defined twice, which could be an issue. +if len(ALL_COMMANDS_DICT) < len(airflow_commands): + seen = set() + dup = [] + for command in airflow_commands: + if command.name not in seen: + seen.add(command.name) + else: + dup.append(command.name) + log.warning( + "The following CLI %d command(s) are defined more than once, " + "CLI behavior when using those will be unpredictable: %s", + len(dup), + dup, + ) + class AirflowHelpFormatter(RichHelpFormatter): """ diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index 28e0eef8c3ff2..c7d00409e9d14 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -20,6 +20,7 @@ import argparse import contextlib +import importlib import io import os import re @@ -29,13 +30,15 @@ from collections import Counter from importlib import reload from pathlib import Path -from unittest.mock import patch +from unittest import mock +from unittest.mock import MagicMock, patch import pytest from airflow.cli import cli_config, cli_parser from airflow.cli.cli_config import ActionCommand, lazy_load_command from airflow.configuration import AIRFLOW_HOME +from airflow.executors.local_executor import LocalExecutor from tests.test_utils.config import conf_vars # Can not be `--snake_case` or contain uppercase letter @@ -133,6 +136,22 @@ def test_subcommand_arg_flag_conflict(self): f"short option flags {conflict_short_option}" ) + @mock.patch.object(LocalExecutor, "get_cli_commands") + def test_dynamic_conflict_detection(self, cli_commands_mock: MagicMock, caplog): + cli_commands_mock.return_value = [ + ActionCommand( + name="webserver", + help="just a command that'll conflict with one defined in core", + func=lambda: None, + args=[], + ) + ] + with caplog.at_level("WARN"): + # force re-evaluation of cli commands (done in top level code) + importlib.reload(cli_parser) + assert len(caplog.messages) == 1 + assert "webserver" in caplog.messages[0] # message mentions the command that's in conflict + def test_falsy_default_value(self): arg = cli_parser.Arg(("--test",), default=0, type=int) parser = argparse.ArgumentParser() From 4c747bb6942f650528a913085ae0093e92d119bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Tue, 15 Aug 2023 16:11:38 -0700 Subject: [PATCH 2/9] rewrite test to not interfer other tests on existing commands --- tests/cli/test_cli_parser.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index c7d00409e9d14..56ff388615dc7 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -36,7 +36,7 @@ import pytest from airflow.cli import cli_config, cli_parser -from airflow.cli.cli_config import ActionCommand, lazy_load_command +from airflow.cli.cli_config import ActionCommand, core_commands, lazy_load_command from airflow.configuration import AIRFLOW_HOME from airflow.executors.local_executor import LocalExecutor from tests.test_utils.config import conf_vars @@ -138,9 +138,17 @@ def test_subcommand_arg_flag_conflict(self): @mock.patch.object(LocalExecutor, "get_cli_commands") def test_dynamic_conflict_detection(self, cli_commands_mock: MagicMock, caplog): + core_commands.append( + ActionCommand( + name="test_command", + help="does nothing", + func=lambda: None, + args=[], + ) + ) cli_commands_mock.return_value = [ ActionCommand( - name="webserver", + name="test_command", help="just a command that'll conflict with one defined in core", func=lambda: None, args=[], @@ -150,7 +158,7 @@ def test_dynamic_conflict_detection(self, cli_commands_mock: MagicMock, caplog): # force re-evaluation of cli commands (done in top level code) importlib.reload(cli_parser) assert len(caplog.messages) == 1 - assert "webserver" in caplog.messages[0] # message mentions the command that's in conflict + assert "test_command" in caplog.messages[0] # message mentions the command that's in conflict def test_falsy_default_value(self): arg = cli_parser.Arg(("--test",), default=0, type=int) From 112d48c391ff6018495b92c512ee8dbdd83ed57a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 16 Aug 2023 12:07:12 -0700 Subject: [PATCH 3/9] throw an exception instead of logging a warning --- airflow/cli/cli_parser.py | 24 ++++++++++-------------- airflow/cli/utils.py | 6 ++++++ tests/cli/test_cli_parser.py | 12 +++++------- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index 489a97f59e616..ee1e60f1b210a 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -24,6 +24,7 @@ from __future__ import annotations import argparse +import collections import logging from argparse import Action from functools import lru_cache @@ -41,11 +42,12 @@ GroupCommand, core_commands, ) +from airflow.cli.utils import CliConflictError from airflow.exceptions import AirflowException from airflow.executors.executor_loader import ExecutorLoader from airflow.utils.helpers import partition -airflow_commands = core_commands +airflow_commands = core_commands.copy() # make a copy to prevent bad interactions in tests log = logging.getLogger(__name__) try: @@ -59,26 +61,20 @@ "a 3.3.0+ version of the Celery provider. If using a Kubernetes executor, install a " "7.4.0+ version of the CNCF provider" ) - # Do no re-raise the exception since we want the CLI to still function for + # Do not re-raise the exception since we want the CLI to still function for # other commands. ALL_COMMANDS_DICT: dict[str, CLICommand] = {sp.name: sp for sp in airflow_commands} + # Check if sub-commands are defined twice, which could be an issue. if len(ALL_COMMANDS_DICT) < len(airflow_commands): - seen = set() - dup = [] - for command in airflow_commands: - if command.name not in seen: - seen.add(command.name) - else: - dup.append(command.name) - log.warning( - "The following CLI %d command(s) are defined more than once, " - "CLI behavior when using those will be unpredictable: %s", - len(dup), - dup, + dup = {k for k, v in collections.Counter([c.name for c in airflow_commands]).items() if v > 1} + raise CliConflictError( + f"The following CLI {len(dup)} command(s) are defined more than once: {sorted(dup)}\n" + f"This can be due to the executor '{ExecutorLoader.get_default_executor_name()}' " + f"redefining core airflow CLI commands." ) diff --git a/airflow/cli/utils.py b/airflow/cli/utils.py index 718d34a6eb75c..a300798ec716c 100644 --- a/airflow/cli/utils.py +++ b/airflow/cli/utils.py @@ -21,6 +21,12 @@ import sys +class CliConflictError(Exception): + """Error for when CLI commands are defined twice by different sources.""" + + pass + + def is_stdout(fileio: io.IOBase) -> bool: """Check whether a file IO is stdout. diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index 56ff388615dc7..f93f70b293b15 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -20,7 +20,6 @@ import argparse import contextlib -import importlib import io import os import re @@ -37,6 +36,7 @@ from airflow.cli import cli_config, cli_parser from airflow.cli.cli_config import ActionCommand, core_commands, lazy_load_command +from airflow.cli.utils import CliConflictError from airflow.configuration import AIRFLOW_HOME from airflow.executors.local_executor import LocalExecutor from tests.test_utils.config import conf_vars @@ -137,7 +137,7 @@ def test_subcommand_arg_flag_conflict(self): ) @mock.patch.object(LocalExecutor, "get_cli_commands") - def test_dynamic_conflict_detection(self, cli_commands_mock: MagicMock, caplog): + def test_dynamic_conflict_detection(self, cli_commands_mock: MagicMock): core_commands.append( ActionCommand( name="test_command", @@ -154,11 +154,9 @@ def test_dynamic_conflict_detection(self, cli_commands_mock: MagicMock, caplog): args=[], ) ] - with caplog.at_level("WARN"): + with pytest.raises(CliConflictError, match="test_command"): # force re-evaluation of cli commands (done in top level code) - importlib.reload(cli_parser) - assert len(caplog.messages) == 1 - assert "test_command" in caplog.messages[0] # message mentions the command that's in conflict + reload(cli_parser) def test_falsy_default_value(self): arg = cli_parser.Arg(("--test",), default=0, type=int) @@ -295,7 +293,7 @@ def test_cli_parser_executors(self, executor, expected_args): parser = cli_parser.get_parser() with pytest.raises(SystemExit) as e: parser.parse_args([expected_arg, "--help"]) - assert e.value.code == 0 + assert e.value.code == 0, stderr.getvalue() stderr = stderr.getvalue() assert "airflow command error" not in stderr From afdfccabc24f6c3ceeff00f0b2260443706467c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 16 Aug 2023 14:26:44 -0700 Subject: [PATCH 4/9] fix test that was passing only because of side effects from other tests --- tests/cli/conftest.py | 17 ++++++++--------- tests/cli/test_cli_parser.py | 15 ++++++++------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/cli/conftest.py b/tests/cli/conftest.py index dcb8f5a9a7ad8..9987afb6833ab 100644 --- a/tests/cli/conftest.py +++ b/tests/cli/conftest.py @@ -23,7 +23,9 @@ from airflow import models from airflow.cli import cli_parser +from airflow.executors import local_executor from airflow.providers.celery.executors import celery_executor, celery_kubernetes_executor +from airflow.providers.cncf.kubernetes.executors import kubernetes_executor, local_kubernetes_executor from tests.test_utils.config import conf_vars # Create custom executors here because conftest is imported first @@ -34,17 +36,14 @@ custom_executor_module.CustomCeleryKubernetesExecutor = type( # type: ignore "CustomCeleryKubernetesExecutor", (celery_kubernetes_executor.CeleryKubernetesExecutor,), {} ) -custom_executor_module.CustomCeleryExecutor = type( # type: ignore - "CustomLocalExecutor", (celery_executor.CeleryExecutor,), {} -) -custom_executor_module.CustomCeleryKubernetesExecutor = type( # type: ignore - "CustomLocalKubernetesExecutor", (celery_kubernetes_executor.CeleryKubernetesExecutor,), {} +custom_executor_module.CustomLocalExecutor = type( # type: ignore + "CustomLocalExecutor", (local_executor.LocalExecutor,), {} ) -custom_executor_module.CustomCeleryExecutor = type( # type: ignore - "CustomKubernetesExecutor", (celery_executor.CeleryExecutor,), {} +custom_executor_module.CustomLocalKubernetesExecutor = type( # type: ignore + "CustomLocalKubernetesExecutor", (local_kubernetes_executor.LocalKubernetesExecutor,), {} ) -custom_executor_module.CustomCeleryKubernetesExecutor = type( # type: ignore - "CustomCeleryKubernetesExecutor", (celery_kubernetes_executor.CeleryKubernetesExecutor,), {} +custom_executor_module.CustomKubernetesExecutor = type( # type: ignore + "CustomKubernetesExecutor", (kubernetes_executor.KubernetesExecutor,), {} ) sys.modules["custom_executor"] = custom_executor_module diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index f93f70b293b15..aa48a15a5437a 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -272,15 +272,16 @@ def test_dag_parser_celery_command_accept_celery_executor(self, executor): [ ("CeleryExecutor", ["celery"]), ("CeleryKubernetesExecutor", ["celery", "kubernetes"]), - ("custom_executor.CustomCeleryExecutor", ["celery"]), - ("custom_executor.CustomCeleryKubernetesExecutor", ["celery", "kubernetes"]), ("KubernetesExecutor", ["kubernetes"]), - ("custom_executor.KubernetesExecutor", ["kubernetes"]), ("LocalExecutor", []), ("LocalKubernetesExecutor", ["kubernetes"]), - ("custom_executor.LocalExecutor", []), - ("custom_executor.LocalKubernetesExecutor", ["kubernetes"]), ("SequentialExecutor", []), + # custom executors are mapped to the regular ones in `conftest.py` + ("custom_executor.CustomLocalExecutor", []), + ("custom_executor.CustomLocalKubernetesExecutor", ["kubernetes"]), + ("custom_executor.CustomCeleryExecutor", ["celery"]), + ("custom_executor.CustomCeleryKubernetesExecutor", ["celery", "kubernetes"]), + ("custom_executor.CustomKubernetesExecutor", ["kubernetes"]), ], ) def test_cli_parser_executors(self, executor, expected_args): @@ -291,9 +292,9 @@ def test_cli_parser_executors(self, executor, expected_args): ) as stderr: reload(cli_parser) parser = cli_parser.get_parser() - with pytest.raises(SystemExit) as e: + with pytest.raises(SystemExit) as e: # running the help command exits, so we prevent that parser.parse_args([expected_arg, "--help"]) - assert e.value.code == 0, stderr.getvalue() + assert e.value.code == 0, stderr.getvalue() # return code 0 == no problem stderr = stderr.getvalue() assert "airflow command error" not in stderr From d1180cd84845e444f83456f8df14b36c6201da19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 16 Aug 2023 14:37:55 -0700 Subject: [PATCH 5/9] improve/remove tests that stopped making sense --- tests/cli/test_cli_parser.py | 36 ++++-------------------------------- 1 file changed, 4 insertions(+), 32 deletions(-) diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index aa48a15a5437a..005544da8f0ab 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -230,42 +230,22 @@ def test_positive_int(self): cli_config.positive_int(allow_zero=True)("-1") @pytest.mark.parametrize( - "executor", + "command", [ "celery", "kubernetes", ], ) - def test_dag_parser_require_celery_executor(self, executor): + def test_executor_specific_commands_not_accessible(self, command): with conf_vars({("core", "executor"): "SequentialExecutor"}), contextlib.redirect_stderr( io.StringIO() ) as stderr: reload(cli_parser) parser = cli_parser.get_parser() with pytest.raises(SystemExit): - parser.parse_args([executor]) + parser.parse_args([command]) stderr = stderr.getvalue() - assert (f"airflow command error: argument GROUP_OR_COMMAND: invalid choice: '{executor}'") in stderr - - @pytest.mark.parametrize( - "executor", - [ - "CeleryExecutor", - "CeleryKubernetesExecutor", - "custom_executor.CustomCeleryExecutor", - "custom_executor.CustomCeleryKubernetesExecutor", - ], - ) - def test_dag_parser_celery_command_accept_celery_executor(self, executor): - with conf_vars({("core", "executor"): executor}), contextlib.redirect_stderr(io.StringIO()) as stderr: - reload(cli_parser) - parser = cli_parser.get_parser() - with pytest.raises(SystemExit): - parser.parse_args(["celery"]) - stderr = stderr.getvalue() - assert ( - "airflow celery command error: the following arguments are required: COMMAND, see help above." - ) in stderr + assert (f"airflow command error: argument GROUP_OR_COMMAND: invalid choice: '{command}'") in stderr @pytest.mark.parametrize( "executor,expected_args", @@ -298,14 +278,6 @@ def test_cli_parser_executors(self, executor, expected_args): stderr = stderr.getvalue() assert "airflow command error" not in stderr - def test_dag_parser_config_command_dont_required_celery_executor(self): - with conf_vars({("core", "executor"): "CeleryExecutor"}), contextlib.redirect_stderr( - io.StringIO() - ) as stdout: - parser = cli_parser.get_parser() - parser.parse_args(["config", "get-value", "celery", "broker-url"]) - assert stdout is not None - def test_non_existing_directory_raises_when_metavar_is_dir_for_db_export_cleaned(self): """Test that the error message is correct when the directory does not exist.""" with contextlib.redirect_stderr(io.StringIO()) as stderr: From 623a72076ed794b40e0500e8d2bd32a79a02637a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 16 Aug 2023 14:51:59 -0700 Subject: [PATCH 6/9] mini fix --- tests/cli/test_cli_parser.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index 005544da8f0ab..6235bc63a58f0 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -29,7 +29,6 @@ from collections import Counter from importlib import reload from pathlib import Path -from unittest import mock from unittest.mock import MagicMock, patch import pytest @@ -136,7 +135,7 @@ def test_subcommand_arg_flag_conflict(self): f"short option flags {conflict_short_option}" ) - @mock.patch.object(LocalExecutor, "get_cli_commands") + @patch.object(LocalExecutor, "get_cli_commands") def test_dynamic_conflict_detection(self, cli_commands_mock: MagicMock): core_commands.append( ActionCommand( From 08d17fcad71b2c0746fb3d723dcd0fd2a831d022 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 16 Aug 2023 15:24:30 -0700 Subject: [PATCH 7/9] fix setup of celery tests --- tests/cli/commands/test_celery_command.py | 24 ++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/tests/cli/commands/test_celery_command.py b/tests/cli/commands/test_celery_command.py index f3e3d391180e6..b97278b17a61c 100644 --- a/tests/cli/commands/test_celery_command.py +++ b/tests/cli/commands/test_celery_command.py @@ -17,6 +17,7 @@ # under the License. from __future__ import annotations +import importlib from argparse import Namespace from tempfile import NamedTemporaryFile from unittest import mock @@ -65,11 +66,12 @@ def test_validate_session_dbapi_exception(self, mock_session): class TestCeleryStopCommand: @classmethod def setup_class(cls): - cls.parser = cli_parser.get_parser() + with conf_vars({("core", "executor"): "CeleryExecutor"}): + importlib.reload(cli_parser) + cls.parser = cli_parser.get_parser() @mock.patch("airflow.cli.commands.celery_command.setup_locations") @mock.patch("airflow.cli.commands.celery_command.psutil.Process") - @conf_vars({("core", "executor"): "CeleryExecutor"}) def test_if_right_pid_is_read(self, mock_process, mock_setup_locations): args = self.parser.parse_args(["celery", "stop"]) pid = "123" @@ -90,7 +92,6 @@ def test_if_right_pid_is_read(self, mock_process, mock_setup_locations): @mock.patch("airflow.cli.commands.celery_command.read_pid_from_pidfile") @mock.patch("airflow.providers.celery.executors.celery_executor.app") @mock.patch("airflow.cli.commands.celery_command.setup_locations") - @conf_vars({("core", "executor"): "CeleryExecutor"}) def test_same_pid_file_is_used_in_start_and_stop( self, mock_setup_locations, mock_celery_app, mock_read_pid_from_pidfile ): @@ -116,7 +117,6 @@ def test_same_pid_file_is_used_in_start_and_stop( @mock.patch("airflow.providers.celery.executors.celery_executor.app") @mock.patch("airflow.cli.commands.celery_command.psutil.Process") @mock.patch("airflow.cli.commands.celery_command.setup_locations") - @conf_vars({("core", "executor"): "CeleryExecutor"}) def test_custom_pid_file_is_used_in_start_and_stop( self, mock_setup_locations, @@ -147,12 +147,13 @@ def test_custom_pid_file_is_used_in_start_and_stop( class TestWorkerStart: @classmethod def setup_class(cls): - cls.parser = cli_parser.get_parser() + with conf_vars({("core", "executor"): "CeleryExecutor"}): + importlib.reload(cli_parser) + cls.parser = cli_parser.get_parser() @mock.patch("airflow.cli.commands.celery_command.setup_locations") @mock.patch("airflow.cli.commands.celery_command.Process") @mock.patch("airflow.providers.celery.executors.celery_executor.app") - @conf_vars({("core", "executor"): "CeleryExecutor"}) def test_worker_started_with_required_arguments(self, mock_celery_app, mock_popen, mock_locations): pid_file = "pid_file" mock_locations.return_value = (pid_file, None, None, None) @@ -208,11 +209,12 @@ def test_worker_started_with_required_arguments(self, mock_celery_app, mock_pope class TestWorkerFailure: @classmethod def setup_class(cls): - cls.parser = cli_parser.get_parser() + with conf_vars({("core", "executor"): "CeleryExecutor"}): + importlib.reload(cli_parser) + cls.parser = cli_parser.get_parser() @mock.patch("airflow.cli.commands.celery_command.Process") @mock.patch("airflow.providers.celery.executors.celery_executor.app") - @conf_vars({("core", "executor"): "CeleryExecutor"}) def test_worker_failure_gracefull_shutdown(self, mock_celery_app, mock_popen): args = self.parser.parse_args(["celery", "worker"]) mock_celery_app.run.side_effect = Exception("Mock exception to trigger runtime error") @@ -226,10 +228,11 @@ def test_worker_failure_gracefull_shutdown(self, mock_celery_app, mock_popen): class TestFlowerCommand: @classmethod def setup_class(cls): - cls.parser = cli_parser.get_parser() + with conf_vars({("core", "executor"): "CeleryExecutor"}): + importlib.reload(cli_parser) + cls.parser = cli_parser.get_parser() @mock.patch("airflow.providers.celery.executors.celery_executor.app") - @conf_vars({("core", "executor"): "CeleryExecutor"}) def test_run_command(self, mock_celery_app): args = self.parser.parse_args( [ @@ -268,7 +271,6 @@ def test_run_command(self, mock_celery_app): @mock.patch("airflow.cli.commands.celery_command.setup_locations") @mock.patch("airflow.cli.commands.celery_command.daemon") @mock.patch("airflow.providers.celery.executors.celery_executor.app") - @conf_vars({("core", "executor"): "CeleryExecutor"}) def test_run_command_daemon(self, mock_celery_app, mock_daemon, mock_setup_locations, mock_pid_file): mock_setup_locations.return_value = ( mock.MagicMock(name="pidfile"), From 923fcdcb50fcaecbff2aea89fa7a450593509df1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 16 Aug 2023 15:57:02 -0700 Subject: [PATCH 8/9] same for kubernetes command tests --- tests/cli/commands/test_kubernetes_command.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/cli/commands/test_kubernetes_command.py b/tests/cli/commands/test_kubernetes_command.py index 1f76220f52e6b..3957790fc9b7d 100644 --- a/tests/cli/commands/test_kubernetes_command.py +++ b/tests/cli/commands/test_kubernetes_command.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import importlib import os import tempfile from unittest import mock @@ -26,12 +27,15 @@ from airflow.cli import cli_parser from airflow.cli.commands import kubernetes_command +from tests.test_utils.config import conf_vars class TestGenerateDagYamlCommand: @classmethod def setup_class(cls): - cls.parser = cli_parser.get_parser() + with conf_vars({("core", "executor"): "KubernetesExecutor"}): + importlib.reload(cli_parser) + cls.parser = cli_parser.get_parser() def test_generate_dag_yaml(self): with tempfile.TemporaryDirectory("airflow_dry_run_test/") as directory: @@ -61,7 +65,9 @@ class TestCleanUpPodsCommand: @classmethod def setup_class(cls): - cls.parser = cli_parser.get_parser() + with conf_vars({("core", "executor"): "KubernetesExecutor"}): + importlib.reload(cli_parser) + cls.parser = cli_parser.get_parser() @mock.patch("kubernetes.client.CoreV1Api.delete_namespaced_pod") @mock.patch("airflow.providers.cncf.kubernetes.kube_client.config.load_incluster_config") From 55a86d11694ab2413cf1edc70c16bcf818ec3549 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Thu, 17 Aug 2023 11:49:48 -0700 Subject: [PATCH 9/9] add little warning about collisions in docstring --- airflow/executors/base_executor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/executors/base_executor.py b/airflow/executors/base_executor.py index 81c441b521f83..5a9c5f30fae68 100644 --- a/airflow/executors/base_executor.py +++ b/airflow/executors/base_executor.py @@ -488,6 +488,7 @@ def get_cli_commands() -> list[GroupCommand]: Override this method to expose commands via Airflow CLI to manage this executor. This can be commands to setup/teardown the executor, inspect state, etc. + Make sure to choose unique names for those commands, to avoid collisions. """ return []