diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index ecf3b18d44bb3..b6528384d1fa0 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -57,11 +57,14 @@ class DefaultHelpParser(argparse.ArgumentParser): def _check_value(self, action, value): """Override _check_value and check conditionally added command""" - executor = conf.get('core', 'EXECUTOR') - if value == 'celery' and executor not in (CELERY_EXECUTOR, CELERY_KUBERNETES_EXECUTOR): - message = f'celery subcommand works only with CeleryExecutor, your current executor: {executor}' - raise ArgumentError(action, message) - if value == 'kubernetes': + if action.dest == 'subcommand' and value == 'celery': + executor = conf.get('core', 'EXECUTOR') + if executor not in (CELERY_EXECUTOR, CELERY_KUBERNETES_EXECUTOR): + message = ( + f'celery subcommand works only with CeleryExecutor, your current executor: {executor}' + ) + raise ArgumentError(action, message) + if action.dest == 'subcommand' and value == 'kubernetes': try: import kubernetes.client # noqa: F401 except ImportError: diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index e6cb9bc82dae0..55ec4927bc0e7 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -25,8 +25,10 @@ from unittest import TestCase import pytest +from parameterized import parameterized from airflow.cli import cli_parser +from tests.test_utils.config import conf_vars # Can not be `--snake_case` or contain uppercase letter ILLEGAL_LONG_OPTION_PATTERN = re.compile("^--[a-z]+_[a-z]+|^--.*[A-Z].*") @@ -193,3 +195,35 @@ def test_positive_int(self): with pytest.raises(argparse.ArgumentTypeError): cli_parser.positive_int(allow_zero=False)('0') cli_parser.positive_int(allow_zero=True)('-1') + + def test_dag_parser_celery_command_require_celery_executor(self): + with conf_vars({('core', 'executor'): 'SequentialExecutor'}), contextlib.redirect_stderr( + io.StringIO() + ) as stderr: + parser = cli_parser.get_parser() + with self.assertRaises(SystemExit): + parser.parse_args(['celery']) + stderr = stderr.getvalue() + assert ( + "airflow command error: argument GROUP_OR_COMMAND: celery subcommand " + "works only with CeleryExecutor, your current executor: SequentialExecutor, see help above." + ) in stderr + + @parameterized.expand(["CeleryExecutor", "CeleryKubernetesExecutor"]) + def test_dag_parser_celery_command_accept_celery_executor(self, executor): + with conf_vars({('core', 'executor'): executor}), contextlib.redirect_stderr(io.StringIO()) as stderr: + parser = cli_parser.get_parser() + with self.assertRaises(SystemExit): + parser.parse_args(['celery']) + stderr = stderr.getvalue() + assert ( + "airflow celery command error: the following arguments are required: COMMAND, see help above." + ) 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