10000 chore: Raise error when using features decorated by working_in_progre… · ecda909/adk-python@8932106 · GitHub
[go: up one dir, main page]
Skip to content

Commit 8932106

Browse files
seanzhougooglecopybara-github
seanzhougoogle
authored and
copybara-github
committed
chore: Raise error when using features decorated by working_in_progress decorator
set environment variable ADK_ALLOW_WIP_FEATURES=true can bypass it. working_in_progress features are not working. ADK users are not supposed to set this environment variable. PiperOrigin-RevId: 771333335
1 parent 233fd20 commit 8932106

File tree

2 files changed

+227
-14
lines changed

2 files changed

+227
-14
lines changed

‎src/google/adk/utils/feature_decorator.py

Lines changed: 50 additions & 5 deletions
+
)
Original file line numberDiff line numberDiff line change
@@ -15,29 +15,52 @@
1515
from __future__ import annotations
1616

1717
import functools
18+
import os
1819
from typing import Callable
1920
from typing import cast
21+
from typing import Optional
2022
from typing import TypeVar
2123
from typing import Union
2224
import warnings
2325

26+
from dotenv import load_dotenv
27+
2428
T = TypeVar("T", bound=Union[Callable, type])
2529

2630

2731
def _make_feature_decorator(
28-
*, label: str, default_message: str
32+
*,
33+
label: str,
34+
default_message: str,
35+
block_usage: bool = False,
36+
bypass_env_var: Optional[str] = None,
2937
) -> Callable[[str], Callable[[T], T]]:
3038
def decorator_factory(message: str = default_message) -> Callable[[T], T]:
3139
def decorator(obj: T) -> T:
3240
obj_name = getattr(obj, "__name__", type(obj).__name__)
33-
warn_msg = f"[{label.upper()}] {obj_name}: {message}"
41+
msg = f"[{label.upper()}] {obj_name}: {message}"
3442

3543
if isinstance(obj, type): # decorating a class
3644
orig_init = obj.__init__
3745

3846
@functools.wraps(orig_init)
3947
def new_init(self, *args, **kwargs):
40-
warnings.warn(warn_msg, category=UserWarning, stacklevel=2)
48+
# Load .env file if dotenv is available
49+
load_dotenv()
50+
51+
# Check if usage should be bypassed via environment variable at call time
52+
should_bypass = (
53+
bypass_env_var is not None
54+
and os.environ.get(bypass_env_var, "").lower() == "true"
55+
)
56+
57+
if should_bypass:
58+
# Bypass completely - no warning, no error
59+
pass
60+
elif block_usage:
61+
raise RuntimeError(msg)
62+
else:
63+
warnings.warn(msg, category=UserWarning, stacklevel=2)
4164
return orig_init(self, *args, **kwargs)
4265

4366
obj.__init__ = new_init # type: ignore[attr-defined]
@@ -47,7 +70,22 @@ def new_init(self, *args, **kwargs):
4770

4871
@functools.wraps(obj)
4972
def wrapper(*args, **kwargs):
50-
warnings.warn(warn_msg, category=UserWarning, stacklevel=2)
73+
# Load .env file if dotenv is available
74+
load_dotenv()
75+
76+
# Check if usage should be bypassed via environment variable at call time
77+
should_bypass = (
78+
bypass_env_var is not None
79+
and os.environ.get(bypass_env_var, "").lower() == "true"
80
81+
82+
if should_bypass:
83+
# Bypass completely - no warning, no error
84+
pass
85+
elif block_usage:
86+
raise RuntimeError(msg)
87+
else:
88+
warnings.warn(msg, category=UserWarning, stacklevel=2)
5189
return obj(*args, **kwargs)
5290

5391
return cast(T, wrapper)
@@ -65,11 +103,18 @@ def wrapper(*args, **kwargs):
65103
working_in_progress = _make_feature_decorator(
66104
label="WIP",
67105
default_message=(
68-
"This feature is a work in progress and may be incomplete or unstable."
106+
"This feature is a work in progress and is not working completely. ADK"
107+
" users are not supposed to use it."
69108
),
109+
block_usage=True,
110+
bypass_env_var="ADK_ALLOW_WIP_FEATURES",
70111
)
71112
"""Mark a class or function as a work in progress.
72113
114+
By default, decorated functions/classes will raise RuntimeError when used.
115+
Set ADK_ALLOW_WIP_FEATURES=true environment variable to bypass this restriction.
116+
ADK users are not supposed to set this environment variable.
117+
73118
Sample usage:
74119
75120
```

‎tests/unittests/utils/test_feature_decorator.py

Lines changed: 177 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import os
2+
import tempfile
13
import warnings
24

35
from google.adk.utils.feature_decorator import experimental
@@ -11,25 +13,176 @@ def run(self):
1113
return "running"
1214

1315

16+
@working_in_progress("function not ready")
17+
def wip_function():
18+
return "executing"
19+
20+
1421
@experimental("api may have breaking change in the future.")
1522
def experimental_fn():
1623
return "executing"
1724

1825

19-
def test_working_in_progress_class_warns():
20-
with warnings.catch_warnings(record=True) as w:
21-
warnings.simplefilter("always")
26+
@experimental("class may change")
27+
class ExperimentalClass:
28+
29+
def run(self):
30+
return "running experimental"
2231

32+
33+
def test_working_in_progress_class_raises_error():
34+
"""Test that WIP class raises RuntimeError by default."""
35+
# Ensure environment variable is not set
36+
if "ADK_ALLOW_WIP_FEATURES" in os.environ:
37+
del os.environ["ADK_ALLOW_WIP_FEATURES"]
38+
39+
try:
2340
feature = IncompleteFeature()
41+
assert False, "Expected RuntimeError to be raised"
42+
except RuntimeError as e:
43+
assert "[WIP] IncompleteFeature:" in str(e)
44+
assert "don't use yet" in str(e)
2445

25-
assert feature.run() == "running"
26-
assert len(w) == 1
27-
assert issubclass(w[0].category, UserWarning)
28-
assert "[WIP] IncompleteFeature:" in str(w[0].message)
29-
assert "don't use yet" in str(w[0].message)
3046

47+
def test_working_in_progress_function_raises_error():
48+
"""Test that WIP function raises RuntimeError by default."""
49+
# Ensure environment variable is not set
50+
if "ADK_ALLOW_WIP_FEATURES" in os.environ:
51+
del os.environ["ADK_ALLOW_WIP_FEATURES"]
52+
53+
try:
54+
result = wip_function()
55+
assert False, "Expected RuntimeError to be raised"
56+
except RuntimeError as e:
57+
assert "[WIP] wip_function:" in str(e)
58+
assert "function not ready" in str(e)
59+
60+
61+
def test_working_in_progress_class_bypassed_with_env_var():
62+
"""Test that WIP class works without warnings when env var is set."""
63+
# Set the bypass environment variable
64+
os.environ["ADK_ALLOW_WIP_FEATURES"] = "true"
65+
66+
try:
67+
with warnings.catch_warnings(record=True) as w:
68+
warnings.simplefilter("always")
69+
70+
feature = IncompleteFeature()
71+
result = feature.run()
72+
73+
assert result == "running"
74+
# Should have no warnings when bypassed
75+
assert len(w) == 0
76+
finally:
77+
# Clean up environment variable
78+
if "ADK_ALLOW_WIP_FEATURES" in os.environ:
79+
del os.environ["ADK_ALLOW_WIP_FEATURES"]
80+
81+
82+
def test_working_in_progress_function_bypassed_with_env_var():
83+
"""Test that WIP function works without warnings when env var is set."""
84+
# Set the bypass environment variable
85+
os.environ["ADK_ALLOW_WIP_FEATURES"] = "true"
86+
87+
try:
88+
with warnings.catch_warnings(record=True) as w:
89+
warnings.simplefilter("always")
90+
91+
result = wip_function()
92+
93+
assert result == "executing"
94+
# Should have no warnings when bypassed
95+
assert len(w) == 0
96+
finally:
97+
# Clean up environment variable
98+
if "ADK_ALLOW_WIP_FEATURES" in os.environ:
99+
del os.environ["ADK_ALLOW_WIP_FEATURES"]
100+
101+
102+
def test_working_in_progress_env_var_case_insensitive():
103+
"""Test that WIP bypass works with different case values."""
104+
test_cases = ["true", "True", "TRUE", "tRuE"]
105+
106+
for case in test_cases:
107+
os.environ["ADK_ALLOW_WIP_FEATURES"] = case
108+
109+
try:
110+
with warnings.catch_warnings(record=True) as w:
111+
warnings.simplefilter("always")
31112

32-
def test_experimental_method_warns():
113+
result = wip_function()
114+
115+
assert result == "executing"
116+
assert len(w) == 0
117+
finally:
118+
if "ADK_ALLOW_WIP_FEATURES" in os.environ:
119+
del os.environ["ADK_ALLOW_WIP_FEATURES"]
120+
121+
122+
def test_working_in_progress_env_var_false_values():
123+
"""Test that WIP still raises errors with false-like env var values."""
124+
false_values = ["false", "False", "FALSE", "0", "", "anything_else"]
125+
126+
for false_val in false_values:
127+
os.environ["ADK_ALLOW_WIP_FEATURES"] = false_val
128+
129+
try:
130+
result = wip_function()
131+
assert False, f"Expected RuntimeError with env var '{false_val}'"
132+
except RuntimeError as e:
133+
assert "[WIP] wip_function:" in str(e)
134+
finally:
135+
if "ADK_ALLOW_WIP_FEATURES" in os.environ:
136+
del os.environ["ADK_ALLOW_WIP_FEATURES"]
137+
138+
139+
def test_working_in_progress_loads_from_dotenv_file():
140+
"""Test that WIP decorator can load environment variables from .env file."""
141+
# Skip test if dotenv is not available
142+
try:
143+
from dotenv import load_dotenv
144+
except ImportError:
145+
import pytest
146+
147+
pytest.skip("python-dotenv not available")
148+
149+
# Ensure environment variable is not set in os.environ
150+
if "ADK_ALLOW_WIP_FEATURES" in os.environ:
151+
del os.environ["ADK_ALLOW_WIP_FEATURES"]
152+
153+
# Create a temporary .env file in current directory
154+
dotenv_path = ".env.test"
155+
156+
try:
157+
# Write the env file
158+
with open(dotenv_path, "w") as f:
159+
f.write("ADK_ALLOW_WIP_FEATURES=true\n")
160+
161+
# Load the environment variables from the file
162+
load_dotenv(dotenv_path)
163+
164+
with warnings.catch_warnings(record=True) as w:
165+
warnings.simplefilter("always")
166+
167+
# This should work because the .env file contains ADK_ALLOW_WIP_FEATURES=true
168+
result = wip_function()
169+
170+
assert result == "executing"
171+
# Should have no warnings when bypassed via .env file
172+
assert len(w) == 0
173+
174+
finally:
175+
# Clean up
176+
try:
177+
os.unlink(dotenv_path)
178+
except FileNotFoundError:
179+
pass
180+
if "ADK_ALLOW_WIP_FEATURES" in os.environ:
181+
del os.environ["ADK_ALLOW_WIP_FEATURES"]
182+
183+
184+
def test_experimental_function_warns():
185+
"""Test that experimental function shows warnings (unchanged behavior)."""
33186
with warnings.catch_warnings(record=True) as w:
34187
warnings.simplefilter("always")
35188

@@ -40,3 +193,18 @@ def test_experimental_method_warns():
40193
assert issubclass(w[0].category, UserWarning)
41194
assert "[EXPERIMENTAL] experimental_fn:" in str(w[0].message)
42195
assert "breaking change in the future" in str(w[0].message)
196+
197+
198+
def test_experimental_class_warns():
199+
"""Test that experimental class shows warnings (unchanged behavior)."""
200+
with warnings.catch_warnings(record=True) as w:
201+
warnings.simplefilter("always")
202+
203+
exp_class = ExperimentalClass()
204+
result = exp_class.run()
205+
206+
assert result == "running experimental"
207+
assert len(w) == 1
208+
assert issubclass(w[0].category, UserWarning)
209+
assert "[EXPERIMENTAL] ExperimentalClass:" in str(w[0].message)
210+
assert "class may change" in str(w[0].message)

0 commit comments

Comments
 (0)
0