feat: Add feature decorator for the feature registry system

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 832503990

tests/unittests/features/test_feature_decorator.py

@@ -0,0 +1,207 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import warnings
+
+from google.adk.features._feature_decorator import experimental
+from google.adk.features._feature_decorator import stable
+from google.adk.features._feature_decorator import working_in_progress
+from google.adk.features._feature_registry import _FEATURE_REGISTRY
+from google.adk.features._feature_registry import _get_feature_config
+from google.adk.features._feature_registry import _register_feature
+from google.adk.features._feature_registry import _WARNED_FEATURES
+from google.adk.features._feature_registry import FeatureConfig
+from google.adk.features._feature_registry import FeatureStage
+import pytest
+
+
+@working_in_progress("WIP_CLASS")
+class IncompleteFeature:
+
+  def run(self):
+    return "running"
+
+
+@working_in_progress("WIP_FUNCTION")
+def wip_function():
+  return "executing"
+
+
+@experimental("EXPERIMENTAL_CLASS")
+class ExperimentalClass:
+
+  def run(self):
+    return "running"
+
+
+@experimental("EXPERIMENTAL_FUNCTION")
+def experimental_function():
+  return "executing"
+
+
+@stable("STABLE_CLASS")
+class StableClass:
+
+  def run(self):
+    return "running"
+
+
+@stable("STABLE_FUNCTION")
+def stable_function():
+  return "executing"
+
+
+@pytest.fixture(autouse=True)
+def reset_env_and_registry(monkeypatch):
+  """Reset environment variables and registry before each test."""
+  # Clean up environment variables
+  for key in list(os.environ.keys()):
+    if key.startswith("ADK_ENABLE_") or key.startswith("ADK_DISABLE_"):
+      monkeypatch.delenv(key, raising=False)
+
+  # Add an existing feature to the registry
+  _register_feature(
+      "ENABLED_EXPERIMENTAL_FEATURE",
+      FeatureConfig(FeatureStage.EXPERIMENTAL, default_on=True),
+  )
+
+  _register_feature(
+      "EXPERIMENTAL_FUNCTION",
+      FeatureConfig(FeatureStage.EXPERIMENTAL, default_on=True),
+  )
+
+
+def test_working_in_progress_stage_mismatch():
+  """Test that working_in_progress is used with a non-WIP stage."""
+  try:
+
+    @working_in_progress("ENABLED_EXPERIMENTAL_FEATURE")
+    def unused_function():  # pylint: disable=unused-variable
+      return "unused"
+
+    assert False, "Expected ValueError to be raised."
+  except ValueError as e:
+    assert (
+        "Feature 'ENABLED_EXPERIMENTAL_FEATURE' is being defined with stage"
+        " 'FeatureStage.WIP', but it was previously registered with stage"
+        " 'FeatureStage.EXPERIMENTAL'."
+        in str(e)
+    )
+
+
+def test_working_in_progress_class_raises_error():
+  """Test that WIP class raises RuntimeError by default."""
+
+  try:
+    IncompleteFeature()
+    assert False, "Expected RuntimeError to be raised."
+  except RuntimeError as e:
+    assert "Feature WIP_CLASS is not enabled." in str(e)
+
+
+def test_working_in_progress_class_bypass_with_env_var(monkeypatch):
+  """Test that WIP class can be bypassed with env var."""
+
+  monkeypatch.setenv("ADK_ENABLE_WIP_CLASS", "true")
+
+  with warnings.catch_warnings(record=True) as w:
+    feature = IncompleteFeature()
+    feature.run()
+    assert len(w) == 1
+    assert "[WIP] feature WIP_CLASS is enabled." in str(w[0].message)
+
+
+def test_working_in_progress_function_raises_error():
+  """Test that WIP function raises RuntimeError by default."""
+
+  try:
+    wip_function()
+    assert False, "Expected RuntimeError to be raised."
+  except RuntimeError as e:
+    assert "Feature WIP_FUNCTION is not enabled." in str(e)
+
+
+def test_working_in_progress_function_bypass_with_env_var(monkeypatch):
+  """Test that WIP function can be bypassed with env var."""
+
+  monkeypatch.setenv("ADK_ENABLE_WIP_FUNCTION", "true")
+
+  with warnings.catch_warnings(record=True) as w:
+    wip_function()
+    assert len(w) == 1
+    assert "[WIP] feature WIP_FUNCTION is enabled." in str(w[0].message)
+
+
+def test_disabled_experimental_class_raises_error():
+  """Test that disabled experimental class raises RuntimeError by default."""
+
+  try:
+    ExperimentalClass()
+    assert False, "Expected RuntimeError to be raised."
+  except RuntimeError as e:
+    assert "Feature EXPERIMENTAL_CLASS is not enabled." in str(e)
+
+
+def test_disabled_experimental_class_bypass_with_env_var(monkeypatch):
+  """Test that disabled experimental class can be bypassed with env var."""
+
+  monkeypatch.setenv("ADK_ENABLE_EXPERIMENTAL_CLASS", "true")
+
+  with warnings.catch_warnings(record=True) as w:
+    feature = ExperimentalClass()
+    feature.run()
+    assert len(w) == 1
+    assert "[EXPERIMENTAL] feature EXPERIMENTAL_CLASS is enabled." in str(
+        w[0].message
+    )
+
+
+def test_enabled_experimental_function_does_not_raise_error():
+  """Test that enabled experimental function does not raise error."""
+
+  with warnings.catch_warnings(record=True) as w:
+    experimental_function()
+    assert len(w) == 1
+    assert "[EXPERIMENTAL] feature EXPERIMENTAL_FUNCTION is enabled." in str(
+        w[0].message
+    )
+
+
+def test_enabled_experimental_function_disabled_by_env_var(monkeypatch):
+  """Test that enabled experimental function can be disabled by env var."""
+
+  monkeypatch.setenv("ADK_DISABLE_EXPERIMENTAL_FUNCTION", "true")
+
+  try:
+    experimental_function()
+    assert False, "Expected RuntimeError to be raised."
+  except RuntimeError as e:
+    assert "Feature EXPERIMENTAL_FUNCTION is not enabled." in str(e)
+
+
+def test_stable_class_does_not_raise_error_or_warn():
+  """Test that stable class does not raise error or warn."""
+
+  with warnings.catch_warnings(record=True) as w:
+    StableClass().run()
+    assert not w
+
+
+def test_stable_function_does_not_raise_error_or_warn():
+  """Test that stable function does not raise error or warn."""
+
+  with warnings.catch_warnings(record=True) as w:
+    stable_function()
+    assert not w

tests/unittests/features/test_feature_registry.py

@@ -17,13 +17,13 @@ from __future__ import annotations
 import os
 import warnings
 
-from google.adk.features.feature_registry import _FEATURE_REGISTRY
-from google.adk.features.feature_registry import _get_feature_config
-from google.adk.features.feature_registry import _register_feature
-from google.adk.features.feature_registry import _WARNED_FEATURES
-from google.adk.features.feature_registry import FeatureConfig
-from google.adk.features.feature_registry import FeatureStage
-from google.adk.features.feature_registry import is_feature_enabled
+from google.adk.features._feature_registry import _FEATURE_REGISTRY
+from google.adk.features._feature_registry import _get_feature_config
+from google.adk.features._feature_registry import _register_feature
+from google.adk.features._feature_registry import _WARNED_FEATURES
+from google.adk.features._feature_registry import FeatureConfig
+from google.adk.features._feature_registry import FeatureStage
+from google.adk.features._feature_registry import is_feature_enabled
 import pytest
 
 FEATURE_CONFIG_WIP = FeatureConfig(FeatureStage.WIP, default_on=False)
@@ -44,17 +44,11 @@ def reset_env_and_registry(monkeypatch):
     if key.startswith("ADK_ENABLE_") or key.startswith("ADK_DISABLE_"):
       monkeypatch.delenv(key, raising=False)
 
-  # Clear registry (but keep it as a dict for adding test entries)
-  _FEATURE_REGISTRY.clear()
-
   # Reset warned features set
   _WARNED_FEATURES.clear()
 
   yield
 
-  # Clean up after test
-  _FEATURE_REGISTRY.clear()
-
   # Reset warned features set
   _WARNED_FEATURES.clear()