refactor(tests): Refactor tests to explicitly handle JSON_SCHEMA_FOR_FUNC_DECL feature flag

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

Author: xuanyang15

tests/unittests/cli/test_cli_feature_options.py

@@ -22,6 +22,7 @@
 from google.adk.features._feature_registry import _WARNED_FEATURES
 from google.adk.features._feature_registry import FeatureName
 from google.adk.features._feature_registry import is_feature_enabled
+from google.adk.features._feature_registry import temporary_feature_override
 import pytest
 
 
@@ -218,17 +219,21 @@ def test_no_enable_features_flag(self):
     """Command works without --enable_features flag."""
     enabled_features = []
 
-    @click.command()
-    @feature_options()
-    def test_cmd():
-      enabled_features.append(
-          is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
-      )
+    with temporary_feature_override(
+        FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False
+    ):
 
-    runner = CliRunner()
-    result = runner.invoke(test_cmd, [], catch_exceptions=False)
-    assert result.exit_code == 0
-    assert enabled_features == [False]
+      @click.command()
+      @feature_options()
+      def test_cmd():
+        enabled_features.append(
+            is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
+        )
+
+      runner = CliRunner()
+      result = runner.invoke(test_cmd, [], catch_exceptions=False)
+      assert result.exit_code == 0
+      assert enabled_features == [False]
 
   def test_preserves_function_metadata(self):
     """Decorator preserves the wrapped function's metadata."""

tests/unittests/tools/application_integration_tool/test_integration_connector_tool.py

@@ -99,36 +99,46 @@ def integration_tool_with_auth(mock_rest_api_tool):
   )
 
 
-def test_get_declaration(integration_tool):
-  """Tests the generation of the function declaration."""
-  declaration = integration_tool._get_declaration()
-
-  assert isinstance(declaration, FunctionDeclaration)
-  assert declaration.name == "test_integration_tool"
-  assert declaration.description == "Test integration tool description."
-
-  # Check parameters schema
-  params = declaration.parameters
-  assert isinstance(params, Schema)
-  print(f"params: {params}")
-  assert params.type == Type.OBJECT
-
-  # Check properties (excluded fields should not be present)
-  assert "user_id" in params.properties
-  assert "connection_name" not in params.properties
-  assert "host" not in params.properties
-  assert "service_name" not in params.properties
-  assert "entity" not in params.properties
-  assert "operation" not in params.properties
-  assert "action" not in params.properties
-  assert "page_size" in params.properties
-  assert "filter" in params.properties
-
-  # Check required fields (optional and excluded fields should not be required)
-  assert "user_id" in params.required
-  assert "page_size" not in params.required
-  assert "filter" not in params.required
-  assert "connection_name" not in params.required
+class TestIntegrationConnectorToolLegacy:
+
+  @pytest.fixture(autouse=True)
+  def disable_feature_flag(self):
+    """Disable the JSON_SCHEMA_FOR_FUNC_DECL feature flag for legacy tests."""
+    with temporary_feature_override(
+        FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False
+    ):
+      yield
+
+  def test_get_declaration(self, integration_tool):
+    """Tests the generation of the function declaration."""
+    declaration = integration_tool._get_declaration()
+
+    assert isinstance(declaration, FunctionDeclaration)
+    assert declaration.name == "test_integration_tool"
+    assert declaration.description == "Test integration tool description."
+
+    # Check parameters schema
+    params = declaration.parameters
+    assert isinstance(params, Schema)
+    print(f"params: {params}")
+    assert params.type == Type.OBJECT
+
+    # Check properties (excluded fields should not be present)
+    assert "user_id" in params.properties
+    assert "connection_name" not in params.properties
+    assert "host" not in params.properties
+    assert "service_name" not in params.properties
+    assert "entity" not in params.properties
+    assert "operation" not in params.properties
+    assert "action" not in params.properties
+    assert "page_size" in params.properties
+    assert "filter" in params.properties
+
+    # Check required fields (optional and excluded fields should not be required)
+    assert "user_id" in params.required
+    assert "page_size" not in params.required
+    assert "filter" not in params.required
+    assert "connection_name" not in params.required
 
 
 @pytest.mark.asyncio
@@ -258,21 +268,27 @@ async def test_run_with_auth_async(
     assert result == {"status": "success", "data": "mock_data"}
 
 
-def test_get_declaration_with_json_schema_feature_enabled(integration_tool):
-  """Tests the generation of the function declaration with JSON schema feature enabled."""
-  with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True):
-    declaration = integration_tool._get_declaration()
+class TestIntegrationConnectorToolWithJsonSchema:
 
-  assert isinstance(declaration, FunctionDeclaration)
-  assert declaration.name == "test_integration_tool"
-  assert declaration.description == "Test integration tool description."
-  assert declaration.parameters is None
-  assert declaration.parameters_json_schema == {
-      "type": "object",
-      "properties": {
-          "user_id": {"type": "string", "description": "User ID"},
-          "page_size": {"type": "integer"},
-          "filter": {"type": "string"},
-      },
-      "required": ["user_id"],
-  }
+  def test_get_declaration_with_json_schema_feature_enabled(
+      self, integration_tool
+  ):
+    """Tests the generation of the function declaration with JSON schema feature enabled."""
+    with temporary_feature_override(
+        FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True
+    ):
+      declaration = integration_tool._get_declaration()
+
+    assert isinstance(declaration, FunctionDeclaration)
+    assert declaration.name == "test_integration_tool"
+    assert declaration.description == "Test integration tool description."
+    assert declaration.parameters is None
+    assert declaration.parameters_json_schema == {
+        "type": "object",
+        "properties": {
+            "user_id": {"type": "string", "description": "User ID"},
+            "page_size": {"type": "integer"},
+            "filter": {"type": "string"},
+        },
+        "required": ["user_id"],
+    }

tests/unittests/tools/mcp_tool/test_mcp_tool.py

@@ -63,66 +63,24 @@ def __init__(
     self.outputSchema = outputSchema
 
 
-class TestMCPTool:
-  """Test suite for MCPTool class."""
+class TestMCPToolLegacy:
+  """Legacy tests for MCPTool."""
+
+  @pytest.fixture(autouse=True)
+  def disable_feature_flag(self):
+    with temporary_feature_override(
+        FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False
+    ):
+      yield
 
   def setup_method(self):
-    """Set up test fixtures."""
     self.mock_mcp_tool = MockMCPTool()
     self.mock_session_manager = Mock(spec=MCPSessionManager)
     self.mock_session = AsyncMock()
     self.mock_session_manager.create_session = AsyncMock(
         return_value=self.mock_session
     )
 
-  def test_init_basic(self):
-    """Test basic initialization without auth."""
-    tool = MCPTool(
-        mcp_tool=self.mock_mcp_tool,
-        mcp_session_manager=self.mock_session_manager,
-    )
-
-    assert tool.name == "test_tool"
-    assert tool.description == "Test tool description"
-    assert tool._mcp_tool == self.mock_mcp_tool
-    assert tool._mcp_session_manager == self.mock_session_manager
-
-  def test_init_with_auth(self):
-    """Test initialization with authentication."""
-    # Create real auth scheme instances instead of mocks
-    from fastapi.openapi.models import OAuth2
-
-    auth_scheme = OAuth2(flows={})
-    auth_credential = AuthCredential(
-        auth_type=AuthCredentialTypes.OAUTH2,
-        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
-    )
-
-    tool = MCPTool(
-        mcp_tool=self.mock_mcp_tool,
-        mcp_session_manager=self.mock_session_manager,
-        auth_scheme=auth_scheme,
-        auth_credential=auth_credential,
-    )
-
-    # The auth config is stored in the parent class _credentials_manager
-    assert tool._credentials_manager is not None
-    assert tool._credentials_manager._auth_config.auth_scheme == auth_scheme
-    assert (
-        tool._credentials_manager._auth_config.raw_auth_credential
-        == auth_credential
-    )
-
-  def test_init_with_empty_description(self):
-    """Test initialization with empty description."""
-    mock_tool = MockMCPTool(description=None)
-    tool = MCPTool(
-        mcp_tool=mock_tool,
-        mcp_session_manager=self.mock_session_manager,
-    )
-
-    assert tool.description == ""
-
   def test_get_declaration(self):
     """Test function declaration generation."""
     tool = MCPTool(
@@ -137,6 +95,25 @@ def test_get_declaration(self):
     assert declaration.description == "Test tool description"
     assert declaration.parameters is not None
 
+
+class TestMCPToolWithJsonSchema:
+  """Tests for MCPTool with JSON_SCHEMA_FOR_FUNC_DECL enabled."""
+
+  @pytest.fixture(autouse=True)
+  def enable_feature_flag(self):
+    with temporary_feature_override(
+        FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True
+    ):
+      yield
+
+  def setup_method(self):
+    self.mock_mcp_tool = MockMCPTool()
+    self.mock_session_manager = Mock(spec=MCPSessionManager)
+    self.mock_session = AsyncMock()
+    self.mock_session_manager.create_session = AsyncMock(
+        return_value=self.mock_session
+    )
+
   def test_get_declaration_with_json_schema_for_func_decl_enabled(self):
     """Test function declaration generation with json schema for func decl enabled."""
     tool = MCPTool(
@@ -202,6 +179,67 @@ def test_get_declaration_with_empty_output_schema_and_json_schema_for_func_decl_
     assert declaration.response is None
     assert not declaration.response_json_schema
 
+
+class TestMCPTool:
+  """Test suite for MCPTool class."""
+
+  def setup_method(self):
+    """Set up test fixtures."""
+    self.mock_mcp_tool = MockMCPTool()
+    self.mock_session_manager = Mock(spec=MCPSessionManager)
+    self.mock_session = AsyncMock()
+    self.mock_session_manager.create_session = AsyncMock(
+        return_value=self.mock_session
+    )
+
+  def test_init_basic(self):
+    """Test basic initialization without auth."""
+    tool = MCPTool(
+        mcp_tool=self.mock_mcp_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+
+    assert tool.name == "test_tool"
+    assert tool.description == "Test tool description"
+    assert tool._mcp_tool == self.mock_mcp_tool
+    assert tool._mcp_session_manager == self.mock_session_manager
+
+  def test_init_with_auth(self):
+    """Test initialization with authentication."""
+    # Create real auth scheme instances instead of mocks
+    from fastapi.openapi.models import OAuth2
+
+    auth_scheme = OAuth2(flows={})
+    auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
+    )
+
+    tool = MCPTool(
+        mcp_tool=self.mock_mcp_tool,
+        mcp_session_manager=self.mock_session_manager,
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
+    )
+
+    # The auth config is stored in the parent class _credentials_manager
+    assert tool._credentials_manager is not None
+    assert tool._credentials_manager._auth_config.auth_scheme == auth_scheme
+    assert (
+        tool._credentials_manager._auth_config.raw_auth_credential
+        == auth_credential
+    )
+
+  def test_init_with_empty_description(self):
+    """Test initialization with empty description."""
+    mock_tool = MockMCPTool(description=None)
+    tool = MCPTool(
+        mcp_tool=mock_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+
+    assert tool.description == ""
+
   @pytest.mark.asyncio
   async def test_run_async_impl_no_auth(self):
     """Test running tool without authentication."""