Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion scripts/prod/common_lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,11 @@ class RestartStrategy(Enum):
ONE_BY_ONE = "one_by_one"
NO_RESTART = "no_restart"

def __str__(self) -> str:
# The accepted CLI token is the value (e.g. "all_at_once"); use it so argparse choices and
# error messages show what the user actually types rather than "RestartStrategy.ALL_AT_ONCE".
return self.value


def restart_strategy_converter(strategy_name: str) -> RestartStrategy:
"""Convert string to RestartStrategy enum with informative error message"""
Expand All @@ -166,8 +171,9 @@ def restart_strategy_converter(strategy_name: str) -> RestartStrategy:
strategy_name = strategy_name.lower()

try:
# Looking an Enum up by value raises ValueError (not KeyError) when no member matches.
return RestartStrategy(strategy_name)
except KeyError:
except ValueError:
valid_strategies = ", ".join([strategy.value for strategy in RestartStrategy])
raise argparse.ArgumentTypeError(
f"Invalid restart strategy '{strategy_name}'. Valid options are: {valid_strategies}"
Expand All @@ -190,6 +196,11 @@ def __init__(self, config_map_name: str, pod_name: str) -> None:
self.config_map_name = config_map_name
self.pod_name = pod_name

def __str__(self) -> str:
# The accepted CLI token is the member name (e.g. "Core"); use it so argparse choices and
# error messages show what the user actually types rather than "Service.Core".
return self.name


class NamespaceAndInstructionArgs:
def __init__(
Expand Down
39 changes: 39 additions & 0 deletions scripts/prod/test_enum_cli_tokens.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
#!/usr/bin/env python3
"""Tests that enum CLI tokens (used in argparse choices/help/errors) match accepted input."""

import argparse

import pytest
from common_lib import RestartStrategy, Service, restart_strategy_converter


def test_restart_strategy_str_is_the_accepted_token():
assert str(RestartStrategy.ALL_AT_ONCE) == "all_at_once"
assert str(RestartStrategy.ONE_BY_ONE) == "one_by_one"
assert str(RestartStrategy.NO_RESTART) == "no_restart"
# The string form round-trips back through the converter.
for strategy in RestartStrategy:
assert restart_strategy_converter(str(strategy)) is strategy


def test_service_str_is_the_accepted_token():
assert str(Service.Core) == "Core"
assert str(Service.SierraCompiler) == "SierraCompiler"
# The string form round-trips back through name lookup (how --service is parsed).
for service in Service:
assert Service[str(service)] is service


def test_invalid_restart_strategy_raises_informative_error():
# Enum value lookup raises ValueError; the converter must translate it to ArgumentTypeError
# with the valid options, rather than letting argparse fall back to a generic message.
with pytest.raises(argparse.ArgumentTypeError) as error_info:
restart_strategy_converter("bogus")
message = str(error_info.value)
assert "all_at_once" in message
assert "one_by_one" in message
assert "no_restart" in message


if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))
6 changes: 4 additions & 2 deletions scripts/prod/update_config_and_restart_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,9 +137,11 @@ def main():
"-j",
"--service",
type=service_type_converter,
choices=list(Service),
default=Service.Core,
help="Service type to operate on; determines configmap and pod names (default: Core)",
metavar="SERVICE",
help="Service type to operate on; determines configmap and pod names. One of: "
+ ", ".join(str(service) for service in Service)
+ " (default: Core)",
)

args_builder.add_argument(
Expand Down
24 changes: 21 additions & 3 deletions scripts/prod/update_config_and_restart_nodes_lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,23 @@
from restarter_lib import ServiceRestarter


class _ConciseHelpFormatter(argparse.RawDescriptionHelpFormatter):
"""Help formatter that shows the metavar once per option.

Renders "-j, --service SERVICE" instead of argparse's default
"-j SERVICE, --service SERVICE". This matches argparse's own behavior in Python 3.13+; this
formatter provides it on the older versions we run on. Inherits RawDescriptionHelpFormatter so
the multi-line usage epilog is still printed verbatim.
"""

def _format_action_invocation(self, action: argparse.Action) -> str:
if not action.option_strings or action.nargs == 0:
return super()._format_action_invocation(action)
default_metavar = self._get_default_metavar_for_optional(action)
args_string = self._format_args(action, default_metavar)
return ", ".join(action.option_strings) + " " + args_string


class ApolloArgsParserBuilder:
"""Builder class for creating argument parsers with required flags and custom arguments."""

Expand All @@ -40,7 +57,7 @@ def __init__(self, description: str, usage_example: str, include_restart_strateg
self.usage_example = usage_example
self.parser = argparse.ArgumentParser(
description=description,
formatter_class=argparse.RawDescriptionHelpFormatter,
formatter_class=_ConciseHelpFormatter,
epilog=usage_example,
)

Expand Down Expand Up @@ -105,9 +122,10 @@ def _add_common_flags(self, include_restart_strategy: bool):
"-t",
"--restart-strategy",
type=restart_strategy_converter,
choices=list(RestartStrategy),
required=True,
help="Strategy for restarting nodes",
metavar="STRATEGY",
help="Strategy for restarting nodes. One of: "
+ ", ".join(str(strategy) for strategy in RestartStrategy),
)

def add_argument(self, *args, **kwargs):
Expand Down
Loading