From c8ab8588a45027c6e3b036a4d376c3cacd8aebf2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= <kamilbregula@apache.org>
Date: Mon, 19 Jul 2021 02:08:40 +0200
Subject: [PATCH 1/3] Improve executor validation in CLI

---
 airflow/cli/cli_parser.py    | 13 ++++++++-----
 tests/cli/test_cli_parser.py | 22 ++++++++++++++++++++++
 2 files changed, 30 insertions(+), 5 deletions(-)

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..7cece6d535e5d 100644
--- a/tests/cli/test_cli_parser.py
+++ b/tests/cli/test_cli_parser.py
@@ -25,10 +25,13 @@
 from unittest import TestCase
 
 import pytest
+from parameterized import parameterized
 
 from airflow.cli import cli_parser
 
 # Can not be `--snake_case` or contain uppercase letter
+from tests.test_utils.config import conf_vars
+
 ILLEGAL_LONG_OPTION_PATTERN = re.compile("^--[a-z]+_[a-z]+|^--.*[A-Z].*")
 # Only can be `-[a-z]` or `-[A-Z]`
 LEGAL_SHORT_OPTION_PATTERN = re.compile("^-[a-zA-z]$")
@@ -193,3 +196,22 @@ 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')
+
+    @parameterized.expand(["CeleryExecutor", "CeleryKubernetesExecutor"])
+    def test_dag_parser_celery_command_required_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

From 374016f245c91368b9d6a06c3e0dca576a635a99 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= <kamilbregula@apache.org>
Date: Mon, 19 Jul 2021 02:12:06 +0200
Subject: [PATCH 2/3] fixup! Improve executor validation in CLI

---
 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 7cece6d535e5d..b643a9e27e8b7 100644
--- a/tests/cli/test_cli_parser.py
+++ b/tests/cli/test_cli_parser.py
@@ -28,10 +28,9 @@
 from parameterized import parameterized
 
 from airflow.cli import cli_parser
-
-# Can not be `--snake_case` or contain uppercase letter
 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].*")
 # Only can be `-[a-z]` or `-[A-Z]`
 LEGAL_SHORT_OPTION_PATTERN = re.compile("^-[a-zA-z]$")

From d8dea8a4288e874167689c078d295117b35c06eb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= <kamilbregula@apache.org>
Date: Mon, 19 Jul 2021 02:26:56 +0200
Subject: [PATCH 3/3] fixup! fixup! Improve executor validation in CLI

---
 tests/cli/test_cli_parser.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py
index b643a9e27e8b7..55ec4927bc0e7 100644
--- a/tests/cli/test_cli_parser.py
+++ b/tests/cli/test_cli_parser.py
@@ -196,8 +196,21 @@ def test_positive_int(self):
             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_required_celery_executor(self, executor):
+    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):
