feat: Add --disable_features CLI option to ADK CLI

This flag can be used to override default feature enable state.

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 858659818
This commit is contained in:
Xuan Yang
2026-01-20 10:45:32 -08:00
committed by Copybara-Service
parent 21f63f66ee
commit 53b67ce634
3 changed files with 179 additions and 33 deletions
+46 -16
View File
@@ -50,28 +50,44 @@ LOG_LEVELS = click.Choice(
)
def _apply_feature_overrides(enable_features: tuple[str, ...]) -> None:
def _apply_feature_overrides(
*,
enable_features: tuple[str, ...] = (),
disable_features: tuple[str, ...] = (),
) -> None:
"""Apply feature overrides from CLI flags.
Args:
enable_features: Tuple of feature names to enable.
disable_features: Tuple of feature names to disable.
"""
feature_overrides: dict[str, bool] = {}
for features_str in enable_features:
for feature_name_str in features_str.split(","):
feature_name_str = feature_name_str.strip()
if not feature_name_str:
continue
try:
feature_name = FeatureName(feature_name_str)
override_feature_enabled(feature_name, True)
except ValueError:
valid_names = ", ".join(f.value for f in FeatureName)
click.secho(
f"WARNING: Unknown feature name '{feature_name_str}'. "
f"Valid names are: {valid_names}",
fg="yellow",
err=True,
)
if feature_name_str:
feature_overrides[feature_name_str] = True
for features_str in disable_features:
for feature_name_str in features_str.split(","):
feature_name_str = feature_name_str.strip()
if feature_name_str:
feature_overrides[feature_name_str] = False
# Apply all overrides
for feature_name_str, enabled in feature_overrides.items():
try:
feature_name = FeatureName(feature_name_str)
override_feature_enabled(feature_name, enabled)
except ValueError:
valid_names = ", ".join(f.value for f in FeatureName)
click.secho(
f"WARNING: Unknown feature name '{feature_name_str}'. "
f"Valid names are: {valid_names}",
fg="yellow",
err=True,
)
def feature_options():
@@ -88,11 +104,25 @@ def feature_options():
),
multiple=True,
)
@click.option(
"--disable_features",
help=(
"Optional. Comma-separated list of feature names to disable. "
"This provides an alternative to environment variables for "
"disabling features. Example: "
"--disable_features=JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING"
),
multiple=True,
)
@functools.wraps(func)
def wrapper(*args, **kwargs):
enable_features = kwargs.pop("enable_features", ())
if enable_features:
_apply_feature_overrides(enable_features)
disable_features = kwargs.pop("disable_features", ())
if enable_features or disable_features:
_apply_feature_overrides(
enable_features=enable_features,
disable_features=disable_features,
)
return func(*args, **kwargs)
return wrapper
+123 -13
View File
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Unit tests for --enable_features CLI option."""
from __future__ import annotations
import click
@@ -42,45 +40,96 @@ class TestApplyFeatureOverrides:
def test_single_feature(self):
"""Single feature name is applied correctly."""
_apply_feature_overrides(("JSON_SCHEMA_FOR_FUNC_DECL",))
_apply_feature_overrides(enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",))
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
def test_comma_separated_features(self):
"""Comma-separated feature names are applied correctly."""
_apply_feature_overrides((
"JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",
))
_apply_feature_overrides(
enable_features=("JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",)
)
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
assert is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
def test_multiple_flag_values(self):
"""Multiple --enable_features flags are applied correctly."""
_apply_feature_overrides((
"JSON_SCHEMA_FOR_FUNC_DECL",
"PROGRESSIVE_SSE_STREAMING",
))
_apply_feature_overrides(
enable_features=(
"JSON_SCHEMA_FOR_FUNC_DECL",
"PROGRESSIVE_SSE_STREAMING",
)
)
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
assert is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
def test_whitespace_handling(self):
"""Whitespace around feature names is stripped."""
_apply_feature_overrides((" JSON_SCHEMA_FOR_FUNC_DECL , COMPUTER_USE ",))
_apply_feature_overrides(
enable_features=(" JSON_SCHEMA_FOR_FUNC_DECL , COMPUTER_USE ",)
)
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
assert is_feature_enabled(FeatureName.COMPUTER_USE)
def test_empty_string_ignored(self):
"""Empty strings in the list are ignored."""
_apply_feature_overrides(("",))
_apply_feature_overrides(enable_features=("",))
# No error should be raised
def test_unknown_feature_warns(self, capsys):
"""Unknown feature names emit a warning."""
_apply_feature_overrides(("UNKNOWN_FEATURE_XYZ",))
_apply_feature_overrides(enable_features=("UNKNOWN_FEATURE_XYZ",))
captured = capsys.readouterr()
assert "WARNING" in captured.err
assert "UNKNOWN_FEATURE_XYZ" in captured.err
assert "Valid names are:" in captured.err
def test_single_disable_feature(self):
"""Single feature name is disabled correctly."""
# First enable a feature
_apply_feature_overrides(enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",))
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
# Then disable it
_apply_feature_overrides(disable_features=("JSON_SCHEMA_FOR_FUNC_DECL",))
assert not is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
def test_comma_separated_disable_features(self):
"""Comma-separated feature names are disabled correctly."""
# First enable features
_apply_feature_overrides(
enable_features=("JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",)
)
# Then disable them
_apply_feature_overrides(
disable_features=(
"JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",
)
)
assert not is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
assert not is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
def test_disable_overrides_enable(self):
"""Disable is applied after enable, so disable wins for same feature."""
_apply_feature_overrides(
enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",),
disable_features=("JSON_SCHEMA_FOR_FUNC_DECL",),
)
# disable_features is processed after enable_features
assert not is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
def test_enable_and_disable_different_features(self):
"""Enable and disable can be used together for different features."""
# First enable a feature that we'll disable
_apply_feature_overrides(enable_features=("PROGRESSIVE_SSE_STREAMING",))
_apply_feature_overrides(
enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",),
disable_features=("PROGRESSIVE_SSE_STREAMING",),
)
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
assert not is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
class TestFeatureOptionsDecorator:
"""Tests for feature_options decorator."""
@@ -195,3 +244,64 @@ class TestFeatureOptionsDecorator:
"my_test_command" in my_test_command.name
or my_test_command.callback.__name__ == "my_test_command"
)
def test_decorator_adds_disable_features_option(self):
"""Decorator adds --disable_features option to command."""
@click.command()
@feature_options()
def test_cmd():
pass
runner = CliRunner()
result = runner.invoke(test_cmd, ["--help"])
assert "--disable_features" in result.output
def test_disable_features_applied_before_command(self):
"""Features are disabled before the command function runs."""
# First enable the feature via override
_apply_feature_overrides(enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",))
feature_was_disabled = []
@click.command()
@feature_options()
def test_cmd():
feature_was_disabled.append(
not is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
)
runner = CliRunner()
runner.invoke(
test_cmd,
["--disable_features=JSON_SCHEMA_FOR_FUNC_DECL"],
catch_exceptions=False,
)
assert feature_was_disabled == [True]
def test_enable_and_disable_together(self):
"""Both --enable_features and --disable_features work together."""
feature_states = []
@click.command()
@feature_options()
def test_cmd():
feature_states.append(
is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
)
feature_states.append(
is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
)
runner = CliRunner()
runner.invoke(
test_cmd,
[
"--enable_features=JSON_SCHEMA_FOR_FUNC_DECL",
"--disable_features=PROGRESSIVE_SSE_STREAMING",
],
catch_exceptions=False,
)
# JSON_SCHEMA_FOR_FUNC_DECL should be enabled
# PROGRESSIVE_SSE_STREAMING should be disabled
assert feature_states == [True, False]
@@ -95,7 +95,10 @@ def test_adk_run():
assert run_command is not None, "Run command not found"
_check_options_in_parameters(
run_command, cli_run.callback, "run", ignore_params={"enable_features"}
run_command,
cli_run.callback,
"run",
ignore_params={"enable_features", "disable_features"},
)
@@ -105,7 +108,10 @@ def test_adk_eval():
assert eval_command is not None, "Eval command not found"
_check_options_in_parameters(
eval_command, cli_eval.callback, "eval", ignore_params={"enable_features"}
eval_command,
cli_eval.callback,
"eval",
ignore_params={"enable_features", "disable_features"},
)
@@ -118,7 +124,7 @@ def test_adk_web():
web_command,
cli_web.callback,
"web",
ignore_params={"verbose", "enable_features"},
ignore_params={"verbose", "enable_features", "disable_features"},
)
@@ -131,7 +137,7 @@ def test_adk_api_server():
api_server_command,
cli_api_server.callback,
"api_server",
ignore_params={"verbose", "enable_features"},
ignore_params={"verbose", "enable_features", "disable_features"},
)