diff --git a/src/google/adk/utils/feature_decorator.py b/src/google/adk/utils/feature_decorator.py index 30163731..e3be0a65 100644 --- a/src/google/adk/utils/feature_decorator.py +++ b/src/google/adk/utils/feature_decorator.py @@ -15,29 +15,52 @@ from __future__ import annotations import functools +import os from typing import Callable from typing import cast +from typing import Optional from typing import TypeVar from typing import Union import warnings +from dotenv import load_dotenv + T = TypeVar("T", bound=Union[Callable, type]) def _make_feature_decorator( - *, label: str, default_message: str + *, + label: str, + default_message: str, + block_usage: bool = False, + bypass_env_var: Optional[str] = None, ) -> Callable[[str], Callable[[T], T]]: def decorator_factory(message: str = default_message) -> Callable[[T], T]: def decorator(obj: T) -> T: obj_name = getattr(obj, "__name__", type(obj).__name__) - warn_msg = f"[{label.upper()}] {obj_name}: {message}" + msg = f"[{label.upper()}] {obj_name}: {message}" if isinstance(obj, type): # decorating a class orig_init = obj.__init__ @functools.wraps(orig_init) def new_init(self, *args, **kwargs): - warnings.warn(warn_msg, category=UserWarning, stacklevel=2) + # Load .env file if dotenv is available + load_dotenv() + + # Check if usage should be bypassed via environment variable at call time + should_bypass = ( + bypass_env_var is not None + and os.environ.get(bypass_env_var, "").lower() == "true" + ) + + if should_bypass: + # Bypass completely - no warning, no error + pass + elif block_usage: + raise RuntimeError(msg) + else: + warnings.warn(msg, category=UserWarning, stacklevel=2) return orig_init(self, *args, **kwargs) obj.__init__ = new_init # type: ignore[attr-defined] @@ -47,7 +70,22 @@ def _make_feature_decorator( @functools.wraps(obj) def wrapper(*args, **kwargs): - warnings.warn(warn_msg, category=UserWarning, stacklevel=2) + # Load .env file if dotenv is available + load_dotenv() + + # Check if usage should be bypassed via environment variable at call time + should_bypass = ( + bypass_env_var is not None + and os.environ.get(bypass_env_var, "").lower() == "true" + ) + + if should_bypass: + # Bypass completely - no warning, no error + pass + elif block_usage: + raise RuntimeError(msg) + else: + warnings.warn(msg, category=UserWarning, stacklevel=2) return obj(*args, **kwargs) return cast(T, wrapper) @@ -65,11 +103,18 @@ def _make_feature_decorator( working_in_progress = _make_feature_decorator( label="WIP", default_message=( - "This feature is a work in progress and may be incomplete or unstable." + "This feature is a work in progress and is not working completely. ADK" + " users are not supposed to use it." ), + block_usage=True, + bypass_env_var="ADK_ALLOW_WIP_FEATURES", ) """Mark a class or function as a work in progress. +By default, decorated functions/classes will raise RuntimeError when used. +Set ADK_ALLOW_WIP_FEATURES=true environment variable to bypass this restriction. +ADK users are not supposed to set this environment variable. + Sample usage: ``` diff --git a/tests/unittests/utils/test_feature_decorator.py b/tests/unittests/utils/test_feature_decorator.py index aa03fc74..eb700ea6 100644 --- a/tests/unittests/utils/test_feature_decorator.py +++ b/tests/unittests/utils/test_feature_decorator.py @@ -1,3 +1,5 @@ +import os +import tempfile import warnings from google.adk.utils.feature_decorator import experimental @@ -11,25 +13,176 @@ class IncompleteFeature: return "running" +@working_in_progress("function not ready") +def wip_function(): + return "executing" + + @experimental("api may have breaking change in the future.") def experimental_fn(): return "executing" -def test_working_in_progress_class_warns(): - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") +@experimental("class may change") +class ExperimentalClass: + def run(self): + return "running experimental" + + +def test_working_in_progress_class_raises_error(): + """Test that WIP class raises RuntimeError by default.""" + # Ensure environment variable is not set + if "ADK_ALLOW_WIP_FEATURES" in os.environ: + del os.environ["ADK_ALLOW_WIP_FEATURES"] + + try: feature = IncompleteFeature() - - assert feature.run() == "running" - assert len(w) == 1 - assert issubclass(w[0].category, UserWarning) - assert "[WIP] IncompleteFeature:" in str(w[0].message) - assert "don't use yet" in str(w[0].message) + assert False, "Expected RuntimeError to be raised" + except RuntimeError as e: + assert "[WIP] IncompleteFeature:" in str(e) + assert "don't use yet" in str(e) -def test_experimental_method_warns(): +def test_working_in_progress_function_raises_error(): + """Test that WIP function raises RuntimeError by default.""" + # Ensure environment variable is not set + if "ADK_ALLOW_WIP_FEATURES" in os.environ: + del os.environ["ADK_ALLOW_WIP_FEATURES"] + + try: + result = wip_function() + assert False, "Expected RuntimeError to be raised" + except RuntimeError as e: + assert "[WIP] wip_function:" in str(e) + assert "function not ready" in str(e) + + +def test_working_in_progress_class_bypassed_with_env_var(): + """Test that WIP class works without warnings when env var is set.""" + # Set the bypass environment variable + os.environ["ADK_ALLOW_WIP_FEATURES"] = "true" + + try: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + feature = IncompleteFeature() + result = feature.run() + + assert result == "running" + # Should have no warnings when bypassed + assert len(w) == 0 + finally: + # Clean up environment variable + if "ADK_ALLOW_WIP_FEATURES" in os.environ: + del os.environ["ADK_ALLOW_WIP_FEATURES"] + + +def test_working_in_progress_function_bypassed_with_env_var(): + """Test that WIP function works without warnings when env var is set.""" + # Set the bypass environment variable + os.environ["ADK_ALLOW_WIP_FEATURES"] = "true" + + try: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + result = wip_function() + + assert result == "executing" + # Should have no warnings when bypassed + assert len(w) == 0 + finally: + # Clean up environment variable + if "ADK_ALLOW_WIP_FEATURES" in os.environ: + del os.environ["ADK_ALLOW_WIP_FEATURES"] + + +def test_working_in_progress_env_var_case_insensitive(): + """Test that WIP bypass works with different case values.""" + test_cases = ["true", "True", "TRUE", "tRuE"] + + for case in test_cases: + os.environ["ADK_ALLOW_WIP_FEATURES"] = case + + try: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + result = wip_function() + + assert result == "executing" + assert len(w) == 0 + finally: + if "ADK_ALLOW_WIP_FEATURES" in os.environ: + del os.environ["ADK_ALLOW_WIP_FEATURES"] + + +def test_working_in_progress_env_var_false_values(): + """Test that WIP still raises errors with false-like env var values.""" + false_values = ["false", "False", "FALSE", "0", "", "anything_else"] + + for false_val in false_values: + os.environ["ADK_ALLOW_WIP_FEATURES"] = false_val + + try: + result = wip_function() + assert False, f"Expected RuntimeError with env var '{false_val}'" + except RuntimeError as e: + assert "[WIP] wip_function:" in str(e) + finally: + if "ADK_ALLOW_WIP_FEATURES" in os.environ: + del os.environ["ADK_ALLOW_WIP_FEATURES"] + + +def test_working_in_progress_loads_from_dotenv_file(): + """Test that WIP decorator can load environment variables from .env file.""" + # Skip test if dotenv is not available + try: + from dotenv import load_dotenv + except ImportError: + import pytest + + pytest.skip("python-dotenv not available") + + # Ensure environment variable is not set in os.environ + if "ADK_ALLOW_WIP_FEATURES" in os.environ: + del os.environ["ADK_ALLOW_WIP_FEATURES"] + + # Create a temporary .env file in current directory + dotenv_path = ".env.test" + + try: + # Write the env file + with open(dotenv_path, "w") as f: + f.write("ADK_ALLOW_WIP_FEATURES=true\n") + + # Load the environment variables from the file + load_dotenv(dotenv_path) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + # This should work because the .env file contains ADK_ALLOW_WIP_FEATURES=true + result = wip_function() + + assert result == "executing" + # Should have no warnings when bypassed via .env file + assert len(w) == 0 + + finally: + # Clean up + try: + os.unlink(dotenv_path) + except FileNotFoundError: + pass + if "ADK_ALLOW_WIP_FEATURES" in os.environ: + del os.environ["ADK_ALLOW_WIP_FEATURES"] + + +def test_experimental_function_warns(): + """Test that experimental function shows warnings (unchanged behavior).""" with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") @@ -40,3 +193,18 @@ def test_experimental_method_warns(): assert issubclass(w[0].category, UserWarning) assert "[EXPERIMENTAL] experimental_fn:" in str(w[0].message) assert "breaking change in the future" in str(w[0].message) + + +def test_experimental_class_warns(): + """Test that experimental class shows warnings (unchanged behavior).""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + exp_class = ExperimentalClass() + result = exp_class.run() + + assert result == "running experimental" + assert len(w) == 1 + assert issubclass(w[0].category, UserWarning) + assert "[EXPERIMENTAL] ExperimentalClass:" in str(w[0].message) + assert "class may change" in str(w[0].message)