chore: import upstream snapshot with attribution

tests/unittests/tools/computer_use/test_computer_use_toolset.py

@@ -0,0 +1,612 @@
+# Copyright 2026 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.
+
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+
+from google.adk.models.llm_request import LlmRequest
+# Use the actual ComputerEnvironment enum from the code
+from google.adk.tools.computer_use.base_computer import BaseComputer
+from google.adk.tools.computer_use.base_computer import ComputerEnvironment
+from google.adk.tools.computer_use.base_computer import ComputerState
+from google.adk.tools.computer_use.computer_use_tool import ComputerUseTool
+from google.adk.tools.computer_use.computer_use_toolset import ComputerUseToolset
+from google.genai import types
+import pytest
+
+
+class MockComputer(BaseComputer):
+  """Mock Computer implementation for testing."""
+
+  def __init__(self):
+    self.initialize_called = False
+    self.close_called = False
+    self._screen_size = (1920, 1080)
+    self._environment = ComputerEnvironment.ENVIRONMENT_BROWSER
+
+  async def initialize(self):
+    self.initialize_called = True
+
+  async def close(self):
+    self.close_called = True
+
+  async def screen_size(self) -> tuple[int, int]:
+    return self._screen_size
+
+  async def environment(self) -> ComputerEnvironment:
+    return self._environment
+
+  # Implement all abstract methods to make this a concrete class
+  async def open_web_browser(self) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def click_at(self, x: int, y: int) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def hover_at(self, x: int, y: int) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def type_text_at(
+      self,
+      x: int,
+      y: int,
+      text: str,
+      press_enter: bool = True,
+      clear_before_typing: bool = True,
+  ) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def scroll_document(self, direction: str) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def scroll_at(
+      self, x: int, y: int, direction: str, magnitude: int
+  ) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def wait(self, seconds: int) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def go_back(self) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def go_forward(self) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def search(self) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def navigate(self, url: str) -> ComputerState:
+    return ComputerState(screenshot=b"test", url=url)
+
+  async def key_combination(self, keys: list[str]) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def drag_and_drop(
+      self, x: int, y: int, destination_x: int, destination_y: int
+  ) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+  async def current_state(self) -> ComputerState:
+    return ComputerState(screenshot=b"test", url="https://example.com")
+
+
+class TestComputerUseToolset:
+  """Test cases for ComputerUseToolset class."""
+
+  @pytest.fixture
+  def mock_computer(self):
+    """Fixture providing a mock computer."""
+    return MockComputer()
+
+  @pytest.fixture
+  def toolset(self, mock_computer):
+    """Fixture providing a ComputerUseToolset instance."""
+    return ComputerUseToolset(computer=mock_computer)
+
+  def test_init(self, mock_computer):
+    """Test ComputerUseToolset initialization."""
+    toolset = ComputerUseToolset(computer=mock_computer)
+
+    assert toolset._computer == mock_computer
+    assert toolset._initialized is False
+
+  @pytest.mark.asyncio
+  async def test_ensure_initialized(self, toolset, mock_computer):
+    """Test that _ensure_initialized calls computer.initialize()."""
+    assert not mock_computer.initialize_called
+    assert not toolset._initialized
+
+    await toolset._ensure_initialized()
+
+    assert mock_computer.initialize_called
+    assert toolset._initialized
+
+  @pytest.mark.asyncio
+  async def test_ensure_initialized_only_once(self, toolset, mock_computer):
+    """Test that _ensure_initialized only calls initialize once."""
+    await toolset._ensure_initialized()
+
+    # Reset the flag to test it's not called again
+    mock_computer.initialize_called = False
+
+    await toolset._ensure_initialized()
+
+    # Should not be called again
+    assert not mock_computer.initialize_called
+    assert toolset._initialized
+
+  @pytest.mark.asyncio
+  async def test_get_tools(self, toolset, mock_computer):
+    """Test that get_tools returns ComputerUseTool instances."""
+    tools = await toolset.get_tools()
+
+    # Should initialize the computer
+    assert mock_computer.initialize_called
+
+    # Should return a list of ComputerUseTool instances
+    assert isinstance(tools, list)
+    assert len(tools) > 0
+    assert all(isinstance(tool, ComputerUseTool) for tool in tools)
+
+    # Each tool should have the correct configuration
+    for tool in tools:
+      assert tool._screen_size == (1920, 1080)
+      # Should use default virtual screen size
+      assert tool._coordinate_space == (1000, 1000)
+
+  @pytest.mark.asyncio
+  async def test_get_tools_excludes_utility_methods(self, toolset):
+    """Test that get_tools excludes utility methods like screen_size, environment, close."""
+    tools = await toolset.get_tools()
+
+    # Get tool function names
+    tool_names = [tool.func.__name__ for tool in tools]
+
+    # Should exclude utility methods
+    excluded_methods = {"screen_size", "environment", "close"}
+    for method in excluded_methods:
+      assert method not in tool_names
+
+    # initialize might be included since it's a concrete method, not just abstract
+    # This is acceptable behavior
+
+    # Should include action methods
+    expected_methods = {
+        "open_web_browser",
+        "click_at",
+        "hover_at",
+        "type_text_at",
+        "scroll_document",
+        "scroll_at",
+        "wait",
+        "go_back",
+        "go_forward",
+        "search",
+        "navigate",
+        "key_combination",
+        "drag_and_drop",
+        "current_state",
+    }
+    for method in expected_methods:
+      assert method in tool_names
+
+  @pytest.mark.asyncio
+  async def test_get_tools_filters_excluded_functions(self, mock_computer):
+    """Test that get_tools filters out excluded functions."""
+    excluded_funcs = ["drag_and_drop", "key_combination"]
+    toolset = ComputerUseToolset(
+        computer=mock_computer,
+        excluded_predefined_functions=excluded_funcs,
+    )
+
+    tools = await toolset.get_tools()
+    tool_names = [tool.func.__name__ for tool in tools]
+
+    for func in excluded_funcs:
+      assert func not in tool_names
+
+    assert "click_at" in tool_names
+
+  @pytest.mark.asyncio
+  async def test_get_tools_with_readonly_context(self, toolset):
+    """Test get_tools with readonly_context parameter."""
+    from google.adk.agents.readonly_context import ReadonlyContext
+
+    readonly_context = MagicMock(spec=ReadonlyContext)
+
+    tools = await toolset.get_tools(readonly_context=readonly_context)
+
+    # Should still return tools (readonly_context doesn't affect behavior currently)
+    assert isinstance(tools, list)
+    assert len(tools) > 0
+
+  @pytest.mark.asyncio
+  async def test_close(self, toolset, mock_computer):
+    """Test that close calls computer.close()."""
+    await toolset.close()
+
+    assert mock_computer.close_called
+
+  @pytest.mark.asyncio
+  async def test_get_tools_creates_tools_with_correct_methods(
+      self, toolset, mock_computer
+  ):
+    """Test that get_tools creates tools with the correct underlying methods."""
+    tools = await toolset.get_tools()
+
+    # Find the click_at tool
+    click_tool = None
+    for tool in tools:
+      if tool.func.__name__ == "click_at":
+        click_tool = tool
+        break
+
+    assert click_tool is not None
+
+    # The tool's function should have the correct name (wrapped method)
+    assert click_tool.func.__name__ == "click_at"
+
+  @pytest.mark.asyncio
+  async def test_get_tools_handles_custom_screen_size(self, mock_computer):
+    """Test get_tools with custom screen size."""
+    mock_computer._screen_size = (2560, 1440)
+
+    toolset = ComputerUseToolset(computer=mock_computer)
+    tools = await toolset.get_tools()
+
+    # All tools should have the custom screen size
+    for tool in tools:
+      assert tool._screen_size == (2560, 1440)
+
+  @pytest.mark.asyncio
+  async def test_get_tools_handles_custom_environment(self, mock_computer):
+    """Test get_tools with custom environment."""
+    mock_computer._environment = ComputerEnvironment.ENVIRONMENT_UNSPECIFIED
+
+    toolset = ComputerUseToolset(computer=mock_computer)
+    tools = await toolset.get_tools()
+
+    # Should still return tools regardless of environment
+    assert isinstance(tools, list)
+    assert len(tools) > 0
+
+  @pytest.mark.asyncio
+  async def test_multiple_get_tools_calls_return_cached_instances(
+      self, toolset
+  ):
+    """Test that multiple get_tools calls return the same cached instances."""
+    tools1 = await toolset.get_tools()
+    tools2 = await toolset.get_tools()
+
+    # Should return the same list instance
+    assert tools1 is tools2
+
+  def test_inheritance(self, toolset):
+    """Test that ComputerUseToolset inherits from BaseToolset."""
+    from google.adk.tools.base_toolset import BaseToolset
+
+    assert isinstance(toolset, BaseToolset)
+
+  @pytest.mark.asyncio
+  async def test_get_tools_method_filtering(self, toolset):
+    """Test that get_tools properly filters methods from BaseComputer."""
+    tools = await toolset.get_tools()
+
+    # Get all method names from the tools
+    tool_method_names = [tool.func.__name__ for tool in tools]
+
+    # Should not include private methods (starting with _)
+    for name in tool_method_names:
+      assert not name.startswith("_")
+
+    # Should not include excluded methods
+    excluded_methods = {"screen_size", "environment", "close"}
+    for excluded in excluded_methods:
+      assert excluded not in tool_method_names
+
+  @pytest.mark.asyncio
+  async def test_computer_method_binding(self, toolset, mock_computer):
+    """Test that tools are properly bound to the computer instance."""
+    tools = await toolset.get_tools()
+
+    # All tools should have wrapped functions with correct names
+    for tool in tools:
+      # Wrapped functions preserve the original method name via functools.wraps
+      assert callable(tool.func)
+      assert not tool.func.__name__.startswith("_")
+
+  @pytest.mark.asyncio
+  async def test_toolset_handles_computer_initialization_failure(
+      self, mock_computer
+  ):
+    """Test that toolset handles computer initialization failure gracefully."""
+
+    # Make initialize raise an exception
+    async def failing_initialize():
+      raise Exception("Initialization failed")
+
+    mock_computer.initialize = failing_initialize
+
+    toolset = ComputerUseToolset(computer=mock_computer)
+
+    # Should raise the exception when trying to get tools
+    with pytest.raises(Exception, match="Initialization failed"):
+      await toolset.get_tools()
+
+  @pytest.mark.asyncio
+  async def test_process_llm_request(self, toolset, mock_computer):
+    """Test that process_llm_request adds tools and computer use configuration."""
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(),
+    )
+
+    await toolset.process_llm_request(
+        tool_context=MagicMock(), llm_request=llm_request
+    )
+
+    # Should add tools to the request
+    assert len(llm_request.tools_dict) > 0
+
+    # Should add computer use configuration
+    assert llm_request.config.tools is not None
+    assert len(llm_request.config.tools) > 0
+
+    # Should have computer use tool
+    computer_use_tools = [
+        tool
+        for tool in llm_request.config.tools
+        if hasattr(tool, "computer_use") and tool.computer_use
+    ]
+    assert len(computer_use_tools) == 1
+
+    # Should have correct environment
+    computer_use_tool = computer_use_tools[0]
+    assert (
+        computer_use_tool.computer_use.environment
+        == types.Environment.ENVIRONMENT_BROWSER
+    )
+
+  @pytest.mark.asyncio
+  async def test_process_llm_request_with_excluded_functions(
+      self, mock_computer
+  ):
+    """Test that process_llm_request passes excluded_predefined_functions."""
+    excluded_funcs = ["drag_and_drop", "key_combination"]
+    toolset = ComputerUseToolset(
+        computer=mock_computer,
+        excluded_predefined_functions=excluded_funcs,
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(),
+    )
+
+    await toolset.process_llm_request(
+        tool_context=MagicMock(), llm_request=llm_request
+    )
+
+    # Should have computer use tool
+    computer_use_tools = [
+        tool
+        for tool in llm_request.config.tools
+        if hasattr(tool, "computer_use") and tool.computer_use
+    ]
+    assert len(computer_use_tools) == 1
+
+    # Should have correct excluded functions
+    computer_use_tool = computer_use_tools[0]
+    assert (
+        computer_use_tool.computer_use.excluded_predefined_functions
+        == excluded_funcs
+    )
+
+  @pytest.mark.asyncio
+  async def test_process_llm_request_with_existing_computer_use(
+      self, toolset, mock_computer
+  ):
+    """Test that process_llm_request doesn't add duplicate computer use configuration."""
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(
+            tools=[
+                types.Tool(
+                    computer_use=types.ComputerUse(
+                        environment=types.Environment.ENVIRONMENT_BROWSER
+                    )
+                )
+            ]
+        ),
+    )
+
+    original_tools_count = len(llm_request.config.tools)
+
+    await toolset.process_llm_request(
+        tool_context=MagicMock(), llm_request=llm_request
+    )
+
+    # Should not add duplicate computer use configuration
+    assert len(llm_request.config.tools) == original_tools_count
+
+    # Should still add the actual tools
+    assert len(llm_request.tools_dict) > 0
+
+  @pytest.mark.asyncio
+  async def test_process_llm_request_error_handling(self, mock_computer):
+    """Test that process_llm_request handles errors gracefully."""
+
+    # Make environment raise an exception
+    async def failing_environment():
+      raise Exception("Environment failed")
+
+    mock_computer.environment = failing_environment
+
+    toolset = ComputerUseToolset(computer=mock_computer)
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(),
+    )
+
+    # Should raise the exception
+    with pytest.raises(Exception, match="Environment failed"):
+      await toolset.process_llm_request(
+          tool_context=MagicMock(), llm_request=llm_request
+      )
+
+  @pytest.mark.asyncio
+  async def test_adapt_computer_use_tool_sync_adapter(self):
+    """Test adapt_computer_use_tool with sync adapter function."""
+    # Create a mock tool
+    mock_func = AsyncMock()
+    original_tool = ComputerUseTool(
+        func=mock_func,
+        screen_size=(1920, 1080),
+        virtual_screen_size=(1000, 1000),
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(),
+    )
+    llm_request.tools_dict["wait"] = original_tool
+
+    # Create a sync adapter function
+    def sync_adapter(original_func):
+      async def adapted_func():
+        return await original_func(5)
+
+      return adapted_func
+
+    # Call the adaptation method
+    await ComputerUseToolset.adapt_computer_use_tool(
+        "wait", sync_adapter, llm_request
+    )
+
+    # Verify the original tool was replaced
+    assert "wait" not in llm_request.tools_dict
+    assert "adapted_func" in llm_request.tools_dict
+
+    # Verify the new tool has correct properties
+    adapted_tool = llm_request.tools_dict["adapted_func"]
+    assert isinstance(adapted_tool, ComputerUseTool)
+    assert adapted_tool._screen_size == (1920, 1080)
+    assert adapted_tool._coordinate_space == (1000, 1000)
+
+  @pytest.mark.asyncio
+  async def test_adapt_computer_use_tool_async_adapter(self):
+    """Test adapt_computer_use_tool with async adapter function."""
+    # Create a mock tool
+    mock_func = AsyncMock()
+    original_tool = ComputerUseTool(
+        func=mock_func,
+        screen_size=(1920, 1080),
+        virtual_screen_size=(1000, 1000),
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(),
+    )
+    llm_request.tools_dict["wait"] = original_tool
+
+    # Create an async adapter function
+    async def async_adapter(original_func):
+      async def adapted_func():
+        return await original_func(5)
+
+      return adapted_func
+
+    # Call the adaptation method
+    await ComputerUseToolset.adapt_computer_use_tool(
+        "wait", async_adapter, llm_request
+    )
+
+    # Verify the original tool was replaced
+    assert "wait" not in llm_request.tools_dict
+    assert "adapted_func" in llm_request.tools_dict
+
+    # Verify the new tool has correct properties
+    adapted_tool = llm_request.tools_dict["adapted_func"]
+    assert isinstance(adapted_tool, ComputerUseTool)
+    assert adapted_tool._screen_size == (1920, 1080)
+    assert adapted_tool._coordinate_space == (1000, 1000)
+
+  @pytest.mark.asyncio
+  async def test_adapt_computer_use_tool_invalid_method(self):
+    """Test adapt_computer_use_tool with invalid method name."""
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(),
+    )
+
+    def adapter(original_func):
+      async def adapted_func():
+        return await original_func()
+
+      return adapted_func
+
+    # Should not raise an exception, just log a warning
+    await ComputerUseToolset.adapt_computer_use_tool(
+        "invalid_method", adapter, llm_request
+    )
+
+    # Should not add any tools
+    assert len(llm_request.tools_dict) == 0
+
+  @pytest.mark.asyncio
+  async def test_adapt_computer_use_tool_excluded_method(self):
+    """Test adapt_computer_use_tool with excluded method name."""
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(),
+    )
+
+    def adapter(original_func):
+      async def adapted_func():
+        return await original_func()
+
+      return adapted_func
+
+    # Should not raise an exception, just log a warning
+    await ComputerUseToolset.adapt_computer_use_tool(
+        "screen_size", adapter, llm_request
+    )
+
+    # Should not add any tools
+    assert len(llm_request.tools_dict) == 0
+
+  @pytest.mark.asyncio
+  async def test_adapt_computer_use_tool_method_not_in_tools_dict(self):
+    """Test adapt_computer_use_tool when method is not in tools_dict."""
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        config=types.GenerateContentConfig(),
+    )
+
+    def adapter(original_func):
+      async def adapted_func():
+        return await original_func()
+
+      return adapted_func
+
+    # Should not raise an exception, just log a warning
+    await ComputerUseToolset.adapt_computer_use_tool(
+        "wait", adapter, llm_request
+    )
+
+    # Should not add any tools
+    assert len(llm_request.tools_dict) == 0