Prevent Dag CLI subcommands from being silently dropped (#72340)

Author: Eason09053360

airflow-core/src/airflow/cli/cli_config.py

@@ -2348,24 +2348,24 @@ class GroupCommand(NamedTuple):
 
 def _remove_dag_id_opt(command: ActionCommand):
     cmd = command._asdict()
-    cmd["args"] = (arg for arg in command.args if arg is not ARG_DAG_ID)
+    cmd["args"] = tuple(arg for arg in command.args if arg is not ARG_DAG_ID)
     return ActionCommand(**cmd)
 
 
+# Subcommands ``DAG.cli()`` exposes, via ``get_parser(dag_parser=True)``.
+DAG_CLI_DAGS_SUBCOMMANDS = ("list-runs", "pause", "unpause", "test")
+DAG_CLI_TASKS_SUBCOMMANDS = ("list", "test")
+
 dag_cli_commands: list[CLICommand] = [
     GroupCommand(
         name="dags",
         help="Manage DAGs",
-        subcommands=[
-            _remove_dag_id_opt(sp)
-            for sp in DAGS_COMMANDS
-            if sp.name in ["backfill", "list-runs", "pause", "unpause", "test"]
-        ],
+        subcommands=[_remove_dag_id_opt(sp) for sp in DAGS_COMMANDS if sp.name in DAG_CLI_DAGS_SUBCOMMANDS],
     ),
     GroupCommand(
         name="tasks",
         help="Manage tasks",
-        subcommands=[_remove_dag_id_opt(sp) for sp in TASKS_COMMANDS if sp.name in ["list", "test", "run"]],
+        subcommands=[_remove_dag_id_opt(sp) for sp in TASKS_COMMANDS if sp.name in DAG_CLI_TASKS_SUBCOMMANDS],
     ),
 ]
 DAG_CLI_DICT: dict[str, CLICommand] = {sp.name: sp for sp in dag_cli_commands}

airflow-core/tests/unit/cli/test_cli_parser.py

@@ -496,6 +496,26 @@ def test_dag_cli_should_display_help(self):
             with pytest.raises(SystemExit):
                 parser.parse_args([*cmd_args, "--help"])
 
+    @pytest.mark.parametrize(
+        ("selected_names", "source_commands"),
+        [
+            pytest.param(cli_config.DAG_CLI_DAGS_SUBCOMMANDS, cli_config.DAGS_COMMANDS, id="dags"),
+            pytest.param(cli_config.DAG_CLI_TASKS_SUBCOMMANDS, cli_config.TASKS_COMMANDS, id="tasks"),
+        ],
+    )
+    def test_dag_cli_subcommands_all_exist(self, selected_names, source_commands):
+        """A name that no longer exists is silently dropped, so guard against stale entries."""
+        assert set(selected_names) <= {command.name for command in source_commands}
+
+    def test_dag_cli_parser_keeps_args_when_rebuilt(self):
+        """``_remove_dag_id_opt`` must not hand argparse a one-shot generator."""
+        cli_parser.get_parser.cache_clear()
+        first = vars(cli_parser.get_parser(dag_parser=True).parse_args(["dags", "pause"]))
+        cli_parser.get_parser.cache_clear()
+        second = vars(cli_parser.get_parser(dag_parser=True).parse_args(["dags", "pause"]))
+        assert "treat_dag_id_as_regex" in first
+        assert first.keys() == second.keys()
+
     def test_positive_int(self):
         assert cli_config.positive_int(allow_zero=True)("1") == 1
         assert cli_config.positive_int(allow_zero=True)("0") == 0