chore: import upstream snapshot with attribution

tests/unittests/integrations/vmaas/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tests/unittests/integrations/vmaas/test_sandbox_client.py

@@ -0,0 +1,364 @@
+# 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.
+
+"""Unit tests for the SandboxClient class."""
+
+import base64
+import json
+import unittest
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.integrations.vmaas.sandbox_client import SandboxClient
+
+
+def _make_response(data: dict) -> MagicMock:
+  """Create a mock HttpResponse with a JSON body."""
+  response = MagicMock()
+  response.body = json.dumps(data)
+  return response
+
+
+class TestSandboxClient(unittest.IsolatedAsyncioTestCase):
+  """Tests for SandboxClient."""
+
+  def setUp(self):
+    """Set up test fixtures."""
+    self.mock_vertexai_client = MagicMock()
+    self.mock_sandbox = MagicMock()
+    self.access_token = "test_token_12345"
+    self.client = SandboxClient(
+        vertexai_client=self.mock_vertexai_client,
+        sandbox=self.mock_sandbox,
+        access_token=self.access_token,
+    )
+
+  def test_init(self):
+    """Test client initialization."""
+    self.assertEqual(self.client._client, self.mock_vertexai_client)
+    self.assertEqual(self.client._sandbox, self.mock_sandbox)
+    self.assertEqual(self.client._access_token, self.access_token)
+
+  def test_update_access_token(self):
+    """Test updating access token."""
+    new_token = "new_token_67890"
+    self.client.update_access_token(new_token)
+    self.assertEqual(self.client._access_token, new_token)
+
+  @patch("asyncio.to_thread")
+  async def test_make_cdp_request(self, mock_to_thread):
+    """Test making a single CDP request."""
+    mock_to_thread.return_value = _make_response({"result": "success"})
+
+    result = await self.client.make_cdp_request(
+        "Page.navigate", {"url": "https://example.com"}
+    )
+
+    self.assertEqual(result, {"result": "success"})
+    mock_to_thread.assert_called_once()
+    call_args = mock_to_thread.call_args
+    # First positional arg is the send_command method
+    self.assertEqual(
+        call_args[0][0],
+        self.mock_vertexai_client.agent_engines.sandboxes.send_command,
+    )
+    # Check keyword args
+    self.assertEqual(call_args[1]["http_method"], "POST")
+    self.assertEqual(call_args[1]["path"], "cdp")
+    self.assertEqual(call_args[1]["access_token"], self.access_token)
+    self.assertEqual(call_args[1]["sandbox_environment"], self.mock_sandbox)
+    self.assertEqual(
+        call_args[1]["request_dict"],
+        {"command": "Page.navigate", "params": {"url": "https://example.com"}},
+    )
+
+  @patch("asyncio.to_thread")
+  async def test_make_cdp_batch_request_with_batch_endpoint(
+      self, mock_to_thread
+  ):
+    """Test making a batch CDP request using batch endpoint."""
+    mock_to_thread.return_value = _make_response(
+        {"results": [{"status": "success"}, {"status": "success"}]}
+    )
+
+    commands = [
+        {
+            "command": "Input.dispatchMouseEvent",
+            "params": {"type": "mousePressed"},
+        },
+        {
+            "command": "Input.dispatchMouseEvent",
+            "params": {"type": "mouseReleased"},
+        },
+    ]
+    result = await self.client.make_cdp_batch_request(commands)
+
+    # Should have 2 results from batch
+    self.assertEqual(len(result), 2)
+    # Should have made 1 call to /cdps
+    self.assertEqual(mock_to_thread.call_count, 1)
+    call_args = mock_to_thread.call_args
+    self.assertEqual(call_args[1]["path"], "cdps")
+
+  @patch("asyncio.to_thread")
+  async def test_make_cdp_batch_request_fallback_sequential(
+      self, mock_to_thread
+  ):
+    """Test batch CDP falls back to sequential when batch endpoint fails."""
+
+    # First call (to /cdps) fails with 404
+    # Subsequent calls (to /cdp) succeed
+    def side_effect(*args, **kwargs):
+      if kwargs.get("path") == "cdps":
+        raise Exception("404 Not Found")
+      return _make_response({"result": "ok"})
+
+    mock_to_thread.side_effect = side_effect
+
+    commands = [
+        {"command": "Command1", "params": {}},
+        {"command": "Command2", "params": {}},
+    ]
+    result = await self.client.make_cdp_batch_request(commands)
+
+    # Should have 2 results (sequential fallback)
+    self.assertEqual(len(result), 2)
+    self.assertEqual(result[0]["status"], "success")
+    self.assertEqual(result[1]["status"], "success")
+    # Should have made 3 calls (1 failed /cdps + 2 sequential /cdp)
+    self.assertEqual(mock_to_thread.call_count, 3)
+
+  @patch("asyncio.to_thread")
+  async def test_get_screenshot(self, mock_to_thread):
+    """Test capturing a screenshot."""
+    # Create a simple PNG-like base64 data
+    png_data = b"\x89PNG\r\n\x1a\n"
+    base64_data = base64.b64encode(png_data).decode()
+
+    mock_to_thread.return_value = _make_response({"data": base64_data})
+
+    result = await self.client.get_screenshot()
+
+    self.assertEqual(result, png_data)
+    call_args = mock_to_thread.call_args
+    self.assertEqual(
+        call_args[1]["request_dict"]["command"], "Page.captureScreenshot"
+    )
+
+  @patch("asyncio.to_thread")
+  async def test_get_current_url(self, mock_to_thread):
+    """Test getting the current URL."""
+    mock_to_thread.return_value = _make_response({
+        "active_tab_id": "tab1",
+        "all_tabs": [
+            {"id": "tab1", "url": "https://example.com", "title": "Example"},
+        ],
+    })
+
+    result = await self.client.get_current_url()
+
+    self.assertEqual(result, "https://example.com")
+    call_args = mock_to_thread.call_args
+    self.assertEqual(call_args[1]["path"], "tabs")
+    self.assertEqual(call_args[1]["http_method"], "GET")
+
+  @patch("asyncio.to_thread")
+  async def test_get_current_url_no_active_tab(self, mock_to_thread):
+    """Test getting URL when no tab is active."""
+    mock_to_thread.return_value = _make_response({
+        "active_tab_id": None,
+        "all_tabs": [],
+    })
+
+    result = await self.client.get_current_url()
+
+    self.assertIsNone(result)
+
+  @patch("asyncio.to_thread")
+  async def test_navigate(self, mock_to_thread):
+    """Test navigating to a URL."""
+    mock_to_thread.return_value = _make_response({"frameId": "frame123"})
+
+    result = await self.client.navigate("https://example.com")
+
+    self.assertEqual(result, {"frameId": "frame123"})
+    call_args = mock_to_thread.call_args
+    self.assertEqual(
+        call_args[1]["request_dict"],
+        {"command": "Page.navigate", "params": {"url": "https://example.com"}},
+    )
+
+  @patch("asyncio.to_thread")
+  async def test_click_at(self, mock_to_thread):
+    """Test clicking at coordinates."""
+    # Return success for batch endpoint
+    mock_to_thread.return_value = _make_response({"results": [{}, {}]})
+
+    await self.client.click_at(100, 200)
+
+    # Should call batch endpoint
+    call_args = mock_to_thread.call_args
+    self.assertEqual(call_args[1]["path"], "cdps")
+    commands = call_args[1]["request_dict"]["commands"]
+    self.assertEqual(len(commands), 2)
+    self.assertEqual(commands[0]["params"]["type"], "mousePressed")
+    self.assertEqual(commands[1]["params"]["type"], "mouseReleased")
+
+  @patch("asyncio.to_thread")
+  async def test_hover_at(self, mock_to_thread):
+    """Test hovering at coordinates."""
+    mock_to_thread.return_value = _make_response({})
+
+    await self.client.hover_at(150, 250)
+
+    call_args = mock_to_thread.call_args
+    self.assertEqual(
+        call_args[1]["request_dict"]["params"],
+        {"type": "mouseMoved", "x": 150, "y": 250},
+    )
+
+  @patch("asyncio.to_thread")
+  async def test_scroll_at_down(self, mock_to_thread):
+    """Test scrolling down."""
+    mock_to_thread.return_value = _make_response({})
+
+    await self.client.scroll_at(100, 200, "down", 300)
+
+    call_args = mock_to_thread.call_args
+    params = call_args[1]["request_dict"]["params"]
+    self.assertEqual(params["type"], "mouseWheel")
+    self.assertEqual(params["x"], 100)
+    self.assertEqual(params["y"], 200)
+    self.assertEqual(params["deltaX"], 0)
+    self.assertEqual(params["deltaY"], 300)  # Positive for down
+
+  @patch("asyncio.to_thread")
+  async def test_scroll_at_up(self, mock_to_thread):
+    """Test scrolling up."""
+    mock_to_thread.return_value = _make_response({})
+
+    await self.client.scroll_at(100, 200, "up", 300)
+
+    call_args = mock_to_thread.call_args
+    params = call_args[1]["request_dict"]["params"]
+    self.assertEqual(params["deltaY"], -300)  # Negative for up
+
+  @patch("asyncio.to_thread")
+  async def test_go_back(self, mock_to_thread):
+    """Test navigating back."""
+    # First call returns navigation history, second navigates
+    mock_to_thread.side_effect = [
+        _make_response({
+            "currentIndex": 1,
+            "entries": [
+                {"id": 1, "url": "https://first.com"},
+                {"id": 2, "url": "https://second.com"},
+            ],
+        }),
+        _make_response({}),  # Navigation response
+    ]
+
+    result = await self.client.go_back()
+
+    self.assertTrue(result)
+    self.assertEqual(mock_to_thread.call_count, 2)
+
+  @patch("asyncio.to_thread")
+  async def test_go_back_at_beginning(self, mock_to_thread):
+    """Test navigating back when at beginning of history."""
+    mock_to_thread.return_value = _make_response({
+        "currentIndex": 0,
+        "entries": [{"id": 1, "url": "https://first.com"}],
+    })
+
+    result = await self.client.go_back()
+
+    self.assertFalse(result)
+    # Should only call once (to get history)
+    self.assertEqual(mock_to_thread.call_count, 1)
+
+  @patch("asyncio.to_thread")
+  async def test_type_text_with_clear_and_enter(self, mock_to_thread):
+    """Test typing text with clear and enter options."""
+    # Return success for batch endpoint
+    mock_to_thread.return_value = _make_response({"results": [{}] * 7})
+
+    await self.client.type_text(
+        "hello", press_enter=True, clear_before_typing=True
+    )
+
+    # Should have: Ctrl+A down, Ctrl+A up, Delete down, Delete up,
+    # insertText, Enter down, Enter up = 7 commands in batch
+    call_args = mock_to_thread.call_args
+    commands = call_args[1]["request_dict"]["commands"]
+    self.assertEqual(len(commands), 7)
+
+  @patch("asyncio.to_thread")
+  async def test_key_combination(self, mock_to_thread):
+    """Test pressing key combinations."""
+    mock_to_thread.return_value = _make_response({"results": [{}] * 4})
+
+    await self.client.key_combination(["control", "c"])
+
+    # Should have: Control down, c down, c up, Control up = 4 commands
+    call_args = mock_to_thread.call_args
+    commands = call_args[1]["request_dict"]["commands"]
+    self.assertEqual(len(commands), 4)
+
+  @patch("asyncio.to_thread")
+  async def test_drag_and_drop(self, mock_to_thread):
+    """Test drag and drop operation."""
+    mock_to_thread.return_value = _make_response({"results": [{}] * 4})
+
+    await self.client.drag_and_drop(10, 20, 100, 200)
+
+    # Should have: mouseMoved (start), mousePressed, mouseMoved (end),
+    # mouseReleased = 4 commands
+    call_args = mock_to_thread.call_args
+    commands = call_args[1]["request_dict"]["commands"]
+    self.assertEqual(len(commands), 4)
+
+  @patch("asyncio.to_thread")
+  async def test_health_check_healthy(self, mock_to_thread):
+    """Test health check when sandbox is healthy."""
+    mock_to_thread.return_value = _make_response({"status": "healthy"})
+
+    result = await self.client.health_check()
+
+    self.assertTrue(result)
+    call_args = mock_to_thread.call_args
+    self.assertEqual(call_args[1]["http_method"], "GET")
+    self.assertEqual(call_args[1]["path"], "")
+
+  @patch("asyncio.to_thread")
+  async def test_health_check_unhealthy(self, mock_to_thread):
+    """Test health check when sandbox is unhealthy."""
+    mock_to_thread.return_value = _make_response({"status": "unhealthy"})
+
+    result = await self.client.health_check()
+
+    self.assertFalse(result)
+
+  @patch("asyncio.to_thread")
+  async def test_health_check_exception(self, mock_to_thread):
+    """Test health check when request fails."""
+    mock_to_thread.side_effect = Exception("Connection failed")
+
+    result = await self.client.health_check()
+
+    self.assertFalse(result)
+
+
+if __name__ == "__main__":
+  unittest.main()

tests/unittests/integrations/vmaas/test_sandbox_computer.py

@@ -0,0 +1,576 @@
+# 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.
+
+"""Unit tests for the AgentEngineSandboxComputer class."""
+
+import time
+import unittest
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.integrations.vmaas.sandbox_computer import _STATE_KEY_ACCESS_TOKEN
+from google.adk.integrations.vmaas.sandbox_computer import _STATE_KEY_AGENT_ENGINE_NAME
+from google.adk.integrations.vmaas.sandbox_computer import _STATE_KEY_SANDBOX_NAME
+from google.adk.integrations.vmaas.sandbox_computer import _STATE_KEY_TOKEN_EXPIRY
+from google.adk.integrations.vmaas.sandbox_computer import AgentEngineSandboxComputer
+from google.adk.tools.computer_use.base_computer import ComputerEnvironment
+from google.adk.tools.computer_use.base_computer import ComputerState
+
+
+class TestAgentEngineSandboxComputer(unittest.IsolatedAsyncioTestCase):
+  """Tests for AgentEngineSandboxComputer."""
+
+  def setUp(self):
+    """Set up test fixtures."""
+    self.project_id = "test-project"
+    self.location = "us-central1"
+    self.service_account = "sa@test-project.iam.gserviceaccount.com"
+
+  def test_init(self):
+    """Test computer initialization."""
+    computer = AgentEngineSandboxComputer(
+        project_id=self.project_id,
+        location=self.location,
+        service_account_email=self.service_account,
+    )
+
+    self.assertEqual(computer._project_id, self.project_id)
+    self.assertEqual(computer._location, self.location)
+    self.assertEqual(computer._service_account_email, self.service_account)
+    self.assertEqual(computer._screen_size, (1280, 720))
+
+  def test_init_with_byos(self):
+    """Test initialization with bring-your-own-sandbox."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    sandbox_name = f"{agent_engine_name}/sandboxEnvironments/456"
+
+    computer = AgentEngineSandboxComputer(
+        project_id=self.project_id,
+        sandbox_name=sandbox_name,
+    )
+
+    # Agent engine name should be extracted from sandbox_name
+    self.assertEqual(computer._agent_engine_name, agent_engine_name)
+    self.assertEqual(computer._sandbox_name, sandbox_name)
+
+  def test_init_with_template_derives_agent_engine(self):
+    """Test agent engine is derived from sandbox_template_name."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    template_name = f"{agent_engine_name}/sandboxEnvironmentTemplates/789"
+
+    computer = AgentEngineSandboxComputer(
+        project_id=self.project_id,
+        sandbox_template_name=template_name,
+    )
+
+    # Agent engine name should be derived from the template name so the
+    # sandbox is created under the template's reasoning engine (no BYOS).
+    self.assertEqual(computer._agent_engine_name, agent_engine_name)
+    self.assertEqual(computer._sandbox_template_name, template_name)
+
+  def test_init_with_snapshot_derives_agent_engine(self):
+    """Test agent engine is derived from sandbox_snapshot_name."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    snapshot_name = f"{agent_engine_name}/sandboxEnvironmentSnapshots/789"
+
+    computer = AgentEngineSandboxComputer(
+        project_id=self.project_id,
+        sandbox_snapshot_name=snapshot_name,
+    )
+
+    self.assertEqual(computer._agent_engine_name, agent_engine_name)
+    self.assertEqual(computer._sandbox_snapshot_name, snapshot_name)
+
+  def test_init_sandbox_name_takes_precedence_over_template(self):
+    """Test sandbox_name wins when both sandbox_name and template are set."""
+    sandbox_engine = (
+        "projects/test/locations/us-central1/reasoningEngines/sandbox"
+    )
+    sandbox_name = f"{sandbox_engine}/sandboxEnvironments/456"
+    template_name = (
+        "projects/test/locations/us-central1/reasoningEngines/template"
+        "/sandboxEnvironmentTemplates/789"
+    )
+
+    computer = AgentEngineSandboxComputer(
+        project_id=self.project_id,
+        sandbox_name=sandbox_name,
+        sandbox_template_name=template_name,
+    )
+
+    self.assertEqual(computer._agent_engine_name, sandbox_engine)
+
+  def test_init_without_sandbox_source_has_no_agent_engine(self):
+    """Test agent engine is None when no sandbox source is provided."""
+    computer = AgentEngineSandboxComputer(project_id=self.project_id)
+    self.assertIsNone(computer._agent_engine_name)
+
+  async def test_ensure_agent_engine_with_template_name(self):
+    """Test _ensure_agent_engine reuses the template's reasoning engine."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    template_name = f"{agent_engine_name}/sandboxEnvironmentTemplates/789"
+    computer = AgentEngineSandboxComputer(sandbox_template_name=template_name)
+    computer._session_state = {}
+
+    result = await computer._ensure_agent_engine()
+
+    self.assertEqual(result, agent_engine_name)
+    # Should not have created or stored a new engine.
+    self.assertNotIn(_STATE_KEY_AGENT_ENGINE_NAME, computer._session_state)
+
+  def test_init_with_vertexai_client(self):
+    """Test initialization with provided vertexai client."""
+    mock_client = MagicMock()
+    computer = AgentEngineSandboxComputer(vertexai_client=mock_client)
+    self.assertEqual(computer._client, mock_client)
+
+  async def test_screen_size(self):
+    """Test screen_size returns hardcoded size."""
+    computer = AgentEngineSandboxComputer()
+    result = await computer.screen_size()
+    self.assertEqual(result, (1280, 720))
+
+  async def test_environment(self):
+    """Test environment returns ENVIRONMENT_BROWSER."""
+    computer = AgentEngineSandboxComputer()
+    result = await computer.environment()
+    self.assertEqual(result, ComputerEnvironment.ENVIRONMENT_BROWSER)
+
+  async def test_ensure_agent_engine_with_sandbox_name(self):
+    """Test _ensure_agent_engine extracts agent engine from sandbox_name."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    sandbox_name = f"{agent_engine_name}/sandboxEnvironments/456"
+    computer = AgentEngineSandboxComputer(sandbox_name=sandbox_name)
+    computer._session_state = {}
+
+    result = await computer._ensure_agent_engine()
+
+    self.assertEqual(result, agent_engine_name)
+    # Should not have touched session state
+    self.assertNotIn(_STATE_KEY_AGENT_ENGINE_NAME, computer._session_state)
+
+  async def test_ensure_agent_engine_from_session_state(self):
+    """Test _ensure_agent_engine uses session state value."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {_STATE_KEY_AGENT_ENGINE_NAME: agent_engine_name}
+
+    result = await computer._ensure_agent_engine()
+
+    self.assertEqual(result, agent_engine_name)
+
+  @patch("google.adk.integrations.vmaas.sandbox_computer.asyncio.to_thread")
+  @patch.object(AgentEngineSandboxComputer, "_get_client")
+  async def test_ensure_agent_engine_creates_new(
+      self, mock_get_client, mock_to_thread
+  ):
+    """Test _ensure_agent_engine creates new agent engine."""
+    new_engine_name = "projects/test/locations/us-central1/reasoningEngines/new"
+
+    mock_client = MagicMock()
+    mock_get_client.return_value = mock_client
+
+    mock_engine = MagicMock()
+    mock_engine.api_resource.name = new_engine_name
+    mock_to_thread.return_value = mock_engine
+
+    computer = AgentEngineSandboxComputer(project_id=self.project_id)
+    computer._session_state = {}
+
+    result = await computer._ensure_agent_engine()
+
+    self.assertEqual(result, new_engine_name)
+    self.assertEqual(
+        computer._session_state[_STATE_KEY_AGENT_ENGINE_NAME], new_engine_name
+    )
+
+  @patch("google.adk.integrations.vmaas.sandbox_computer.asyncio.to_thread")
+  @patch.object(AgentEngineSandboxComputer, "_get_client")
+  async def test_get_sandbox_with_constructor_value(
+      self, mock_get_client, mock_to_thread
+  ):
+    """Test _get_sandbox uses constructor value (BYOS mode)."""
+    sandbox_name = "projects/test/sandboxEnvironments/123"
+
+    mock_sandbox = MagicMock()
+    mock_sandbox.name = sandbox_name
+    mock_to_thread.return_value = mock_sandbox
+
+    mock_client = MagicMock()
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer(sandbox_name=sandbox_name)
+    computer._session_state = {}
+
+    result_name, result_sandbox = await computer._get_sandbox()
+
+    self.assertEqual(result_name, sandbox_name)
+    self.assertEqual(result_sandbox, mock_sandbox)
+
+  @patch("google.adk.integrations.vmaas.sandbox_computer.asyncio.to_thread")
+  @patch.object(AgentEngineSandboxComputer, "_get_client")
+  async def test_get_sandbox_from_session_state(
+      self, mock_get_client, mock_to_thread
+  ):
+    """Test _get_sandbox uses session state value."""
+    sandbox_name = "projects/test/sandboxEnvironments/123"
+
+    mock_sandbox = MagicMock()
+    mock_to_thread.return_value = mock_sandbox
+
+    mock_client = MagicMock()
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {_STATE_KEY_SANDBOX_NAME: sandbox_name}
+
+    result_name, result_sandbox = await computer._get_sandbox()
+
+    self.assertEqual(result_name, sandbox_name)
+    self.assertEqual(result_sandbox, mock_sandbox)
+
+  async def test_get_access_token_cached(self):
+    """Test _get_access_token uses cached token."""
+    sandbox_name = "projects/test/sandboxEnvironments/123"
+    cached_token = "cached_token_123"
+    # Set expiry far in the future
+    token_expiry = time.time() + 3600
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {
+        _STATE_KEY_ACCESS_TOKEN: cached_token,
+        _STATE_KEY_TOKEN_EXPIRY: token_expiry,
+    }
+
+    result = await computer._get_access_token(sandbox_name)
+
+    self.assertEqual(result, cached_token)
+
+  @patch("google.adk.integrations.vmaas.sandbox_computer.asyncio.to_thread")
+  @patch.object(AgentEngineSandboxComputer, "_get_client")
+  async def test_get_access_token_generates_new_when_expired(
+      self, mock_get_client, mock_to_thread
+  ):
+    """Test _get_access_token generates new token when expired."""
+    sandbox_name = "projects/test/sandboxEnvironments/123"
+    new_token = "new_token_456"
+    # Set expiry in the past
+    token_expiry = time.time() - 100
+
+    mock_to_thread.return_value = new_token
+    mock_client = MagicMock()
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer(
+        service_account_email=self.service_account
+    )
+    computer._session_state = {
+        _STATE_KEY_ACCESS_TOKEN: "old_token",
+        _STATE_KEY_TOKEN_EXPIRY: token_expiry,
+    }
+
+    result = await computer._get_access_token(sandbox_name)
+
+    self.assertEqual(result, new_token)
+    self.assertEqual(
+        computer._session_state[_STATE_KEY_ACCESS_TOKEN], new_token
+    )
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_click_at(self, mock_get_client):
+    """Test click_at method."""
+    mock_client = AsyncMock()
+    mock_client.click_at = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.click_at(100, 200)
+
+    mock_client.click_at.assert_called_once_with(100, 200)
+    self.assertIsInstance(result, ComputerState)
+    self.assertEqual(result.screenshot, b"png_data")
+    self.assertEqual(result.url, "https://example.com")
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_hover_at(self, mock_get_client):
+    """Test hover_at method."""
+    mock_client = AsyncMock()
+    mock_client.hover_at = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.hover_at(150, 250)
+
+    mock_client.hover_at.assert_called_once_with(150, 250)
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_type_text_at(self, mock_get_client):
+    """Test type_text_at method."""
+    mock_client = AsyncMock()
+    mock_client.type_text_at = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.type_text_at(
+        100,
+        200,
+        "hello",
+        press_enter=True,
+        clear_before_typing=False,
+    )
+
+    mock_client.type_text_at.assert_called_once_with(
+        x=100,
+        y=200,
+        text="hello",
+        press_enter=True,
+        clear_before_typing=False,
+    )
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_scroll_document(self, mock_get_client):
+    """Test scroll_document method."""
+    mock_client = AsyncMock()
+    mock_client.scroll_at = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.scroll_document("down")
+
+    # Should scroll at center of screen
+    mock_client.scroll_at.assert_called_once_with(640, 360, "down", 400)
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_scroll_at(self, mock_get_client):
+    """Test scroll_at method."""
+    mock_client = AsyncMock()
+    mock_client.scroll_at = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.scroll_at(100, 200, "up", 500)
+
+    mock_client.scroll_at.assert_called_once_with(100, 200, "up", 500)
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_navigate(self, mock_get_client):
+    """Test navigate method."""
+    mock_client = AsyncMock()
+    mock_client.navigate = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://newsite.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.navigate("https://newsite.com")
+
+    mock_client.navigate.assert_called_once_with("https://newsite.com")
+    self.assertIsInstance(result, ComputerState)
+    self.assertEqual(result.url, "https://newsite.com")
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_search(self, mock_get_client):
+    """Test search method navigates to search engine."""
+    mock_client = AsyncMock()
+    mock_client.navigate = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(
+        return_value="https://www.google.com"
+    )
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer(
+        search_engine_url="https://www.google.com"
+    )
+    computer._session_state = {}
+
+    result = await computer.search()
+
+    mock_client.navigate.assert_called_once_with("https://www.google.com")
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_go_back(self, mock_get_client):
+    """Test go_back method."""
+    mock_client = AsyncMock()
+    mock_client.go_back = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://prev.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.go_back()
+
+    mock_client.go_back.assert_called_once()
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_go_forward(self, mock_get_client):
+    """Test go_forward method."""
+    mock_client = AsyncMock()
+    mock_client.go_forward = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://next.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.go_forward()
+
+    mock_client.go_forward.assert_called_once()
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_key_combination(self, mock_get_client):
+    """Test key_combination method."""
+    mock_client = AsyncMock()
+    mock_client.key_combination = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.key_combination(["control", "c"])
+
+    mock_client.key_combination.assert_called_once_with(["control", "c"])
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_drag_and_drop(self, mock_get_client):
+    """Test drag_and_drop method."""
+    mock_client = AsyncMock()
+    mock_client.drag_and_drop = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.drag_and_drop(10, 20, 100, 200)
+
+    mock_client.drag_and_drop.assert_called_once_with(10, 20, 100, 200)
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_wait(self, mock_get_client):
+    """Test wait method."""
+    mock_client = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    # Use a short wait time for testing
+    start_time = time.time()
+    result = await computer.wait(1)
+    elapsed = time.time() - start_time
+
+    self.assertGreaterEqual(elapsed, 0.9)  # Allow some margin
+    self.assertIsInstance(result, ComputerState)
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_current_state(self, mock_get_client):
+    """Test current_state method."""
+    mock_client = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="https://example.com")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.current_state()
+
+    self.assertIsInstance(result, ComputerState)
+    self.assertEqual(result.screenshot, b"png_data")
+    self.assertEqual(result.url, "https://example.com")
+
+  @patch.object(AgentEngineSandboxComputer, "_get_sandbox_client")
+  async def test_open_web_browser(self, mock_get_client):
+    """Test open_web_browser method returns current state."""
+    mock_client = AsyncMock()
+    mock_client.get_screenshot = AsyncMock(return_value=b"png_data")
+    mock_client.get_current_url = AsyncMock(return_value="about:blank")
+    mock_get_client.return_value = mock_client
+
+    computer = AgentEngineSandboxComputer()
+    computer._session_state = {}
+
+    result = await computer.open_web_browser()
+
+    # open_web_browser is a no-op for sandbox, just returns current state
+    self.assertIsInstance(result, ComputerState)
+
+  async def test_initialize_is_noop(self):
+    """Test initialize does nothing (lazy provisioning)."""
+    computer = AgentEngineSandboxComputer()
+    # Should not raise
+    await computer.initialize()
+
+  async def test_close_is_noop(self):
+    """Test close does nothing (TTL-based cleanup)."""
+    computer = AgentEngineSandboxComputer()
+    # Should not raise
+    await computer.close()
+
+
+if __name__ == "__main__":
+  unittest.main()