1365 lines
46 KiB
Python
1365 lines
46 KiB
Python
# 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.
|
|
|
|
import os
|
|
|
|
os.environ["ARTEMIS_USE_FILE_API"] = "false"
|
|
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
import json
|
|
from google.genai import types
|
|
|
|
from artemis.agents.explorer.explorer import Explorer
|
|
from artemis.context import ArtemisContext
|
|
from artemis.graph.state import State
|
|
|
|
MOCK_PROMPT_JSON = """{
|
|
"IDENTITY": "You are the agentic UI Explorer designed to identify objects, scan text, and extract UI attributes on a phone screen.",
|
|
"OPERATING PRINCIPLES": [
|
|
"Confidence-Driven Cost-Benefit Strategy: Maximize localization accuracy...",
|
|
"Tool Capabilities: You have three classes of tools...",
|
|
"Parallel vs Sequential Invocations: You are fully supported..."
|
|
]
|
|
}"""
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def disable_file_api(monkeypatch):
|
|
monkeypatch.setenv("ARTEMIS_USE_FILE_API", "false")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_run():
|
|
# Mock context and state
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/test_screenshot.jpg"
|
|
|
|
# Mock Gemini Client and Response
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.text = ""
|
|
mock_func_call = MagicMock()
|
|
mock_func_call.name = "submit_answer"
|
|
mock_func_call.args = {
|
|
"candidates": [
|
|
{
|
|
"label": "S1",
|
|
"coords": [500, 500],
|
|
"description": "Target button",
|
|
}
|
|
],
|
|
"fallback_message": "",
|
|
}
|
|
mock_response.function_calls = [mock_func_call]
|
|
|
|
mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
|
|
|
|
# Mock storage manager
|
|
mock_storage = MagicMock()
|
|
mock_storage.get_image.return_value = "dummy_hash"
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
# Custom mock open to handle both binary and text reads
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
# Provide a sequence of bytes for iter(lambda: f.read(4096), b"") and b64 encode
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_data",
|
|
b"",
|
|
b"fake_image_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
result = await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
screenshot_path="/tmp/test_screenshot.jpg",
|
|
minimal_list="",
|
|
state=mock_state,
|
|
)
|
|
|
|
expected_outcome = json.dumps(mock_func_call.args, ensure_ascii=False)
|
|
assert result == expected_outcome
|
|
assert mock_client.aio.models.generate_content.called
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_submit_answer():
|
|
# Mock context and state
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/test_screenshot.jpg"
|
|
|
|
# Mock Gemini Client and Function Call for submit_answer
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.text = ""
|
|
|
|
mock_func_call = MagicMock()
|
|
mock_func_call.name = "submit_answer"
|
|
mock_func_call.args = {
|
|
"candidates": [
|
|
{
|
|
"label": "S1",
|
|
"coords": [500, 500],
|
|
"description": "Target button",
|
|
}
|
|
],
|
|
"fallback_message": "",
|
|
}
|
|
mock_response.function_calls = [mock_func_call]
|
|
|
|
mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
|
|
|
|
# Mock storage manager
|
|
mock_storage = MagicMock()
|
|
mock_storage.get_image.return_value = "dummy_hash"
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_data",
|
|
b"",
|
|
b"fake_image_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
result = await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
screenshot_path="/tmp/test_screenshot.jpg",
|
|
minimal_list="",
|
|
state=mock_state,
|
|
)
|
|
|
|
# Ensure it returned the structured JSON outcome
|
|
expected_outcome = json.dumps(mock_func_call.args, ensure_ascii=False)
|
|
assert result == expected_outcome
|
|
assert mock_client.aio.models.generate_content.called
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_submit_answer_self_correction():
|
|
# Mock context and state
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/test_screenshot.jpg"
|
|
|
|
# Response 1: Returns invalid coords [1500, -200] (out of bounds)
|
|
mock_response1 = MagicMock()
|
|
mock_response1.text = ""
|
|
|
|
mock_func_call1 = MagicMock()
|
|
mock_func_call1.name = "submit_answer"
|
|
mock_func_call1.args = {
|
|
"candidates": [
|
|
{
|
|
"label": "S1",
|
|
"coords": [1500, -200],
|
|
"description": "Ambiguous button",
|
|
}
|
|
]
|
|
}
|
|
|
|
mock_response1.candidates = [
|
|
MagicMock()
|
|
] # Provide dummy candidate content for loop history tracking
|
|
mock_response1.candidates[0].content = types.Content(
|
|
role="model",
|
|
parts=[
|
|
types.Part.from_function_call(
|
|
name="submit_answer",
|
|
args={
|
|
"candidates": [
|
|
{
|
|
"label": "S1",
|
|
"coords": [1500, -200],
|
|
"description": "Ambiguous button",
|
|
}
|
|
]
|
|
},
|
|
)
|
|
],
|
|
)
|
|
mock_response1.function_calls = [mock_func_call1]
|
|
|
|
# Response 2: Returns corrected valid coords [500, 500]
|
|
mock_response2 = MagicMock()
|
|
mock_response2.text = ""
|
|
|
|
mock_func_call2 = MagicMock()
|
|
mock_func_call2.name = "submit_answer"
|
|
mock_func_call2.args = {
|
|
"candidates": [
|
|
{
|
|
"label": "S1",
|
|
"coords": [500, 500],
|
|
"description": "Corrected button",
|
|
}
|
|
]
|
|
}
|
|
|
|
mock_response2.candidates = [MagicMock()]
|
|
mock_response2.candidates[0].content = types.Content(
|
|
role="model",
|
|
parts=[
|
|
types.Part.from_function_call(
|
|
name="submit_answer",
|
|
args={
|
|
"candidates": [
|
|
{
|
|
"label": "S1",
|
|
"coords": [500, 500],
|
|
"description": "Corrected button",
|
|
}
|
|
]
|
|
},
|
|
)
|
|
],
|
|
)
|
|
mock_response2.function_calls = [mock_func_call2]
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_client.aio.models.generate_content = AsyncMock(
|
|
side_effect=[mock_response1, mock_response2]
|
|
)
|
|
|
|
# Mock storage manager
|
|
mock_storage = MagicMock()
|
|
mock_storage.get_image.return_value = "dummy_hash"
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_data",
|
|
b"",
|
|
b"fake_image_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
result = await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
screenshot_path="/tmp/test_screenshot.jpg",
|
|
minimal_list="",
|
|
state=mock_state,
|
|
)
|
|
|
|
# Assert loops successfully completed on the second corrected turn
|
|
expected_args = {
|
|
"candidates": [
|
|
{
|
|
"label": "S1",
|
|
"coords": [500, 500],
|
|
"description": "Corrected button",
|
|
}
|
|
]
|
|
}
|
|
assert result == json.dumps(expected_args, ensure_ascii=False)
|
|
assert mock_client.aio.models.generate_content.call_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_initial_visual_marking():
|
|
# Mock context and state with latest_ui_hierarchy
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/test_screenshot.jpg"
|
|
mock_state.latest_ui_hierarchy = [{"text": "Clickable Button", "bounds": "[100,100][200,200]"}]
|
|
|
|
# Mock Gemini Client and response
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.text = ""
|
|
|
|
mock_func_call = MagicMock()
|
|
mock_func_call.name = "submit_answer"
|
|
mock_func_call.args = {
|
|
"candidates": [{"label": "1", "coords": [500, 500], "description": "Visual target"}]
|
|
}
|
|
mock_response.function_calls = [mock_func_call]
|
|
mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
|
|
|
|
# Mock storage manager
|
|
mock_storage = MagicMock()
|
|
mock_storage.get_image.return_value = "dummy_hash"
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_marked_data",
|
|
b"",
|
|
b"fake_image_marked_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots,
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
screenshot_path="/tmp/test_screenshot.jpg",
|
|
state=mock_state,
|
|
minimal_list="",
|
|
)
|
|
|
|
# Ensure draw_dots was successfully called with calculated center coordinates of bounds [100,100][200,200] -> [150, 150]
|
|
mock_draw_dots.assert_called_once()
|
|
points_arg = mock_draw_dots.call_args[0][1]
|
|
assert points_arg == [[150, 150]]
|
|
|
|
labels_arg = mock_draw_dots.call_args[0][2]
|
|
assert labels_arg == ["1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_initial_visual_marking_previous_screenshot():
|
|
from artemis.data_engine.models import ImageRecord
|
|
|
|
# Mock context and state
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
# current screenshot is different
|
|
mock_state.latest_screenshot = "/tmp/current_screenshot.jpg"
|
|
mock_state.latest_ui_hierarchy = [{"text": "Current Button", "bounds": "[300,300][400,400]"}]
|
|
|
|
# Mock Gemini Client and response
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.text = ""
|
|
|
|
mock_func_call = MagicMock()
|
|
mock_func_call.name = "submit_answer"
|
|
mock_func_call.args = {
|
|
"candidates": [{"label": "1", "coords": [500, 500], "description": "Visual target"}]
|
|
}
|
|
mock_response.function_calls = [mock_func_call]
|
|
mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
|
|
|
|
# Mock storage manager to return correct previous image record from DB
|
|
mock_storage = MagicMock()
|
|
|
|
# Mock ImageRecord with the previous screenshot's UI hierarchy
|
|
mock_record = MagicMock(spec=ImageRecord)
|
|
mock_record.ui_tree = [{"text": "Previous Button", "bounds": "[100,100][200,200]"}]
|
|
mock_storage.get_image.return_value = mock_record
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_marked_data",
|
|
b"",
|
|
b"fake_image_marked_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots,
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
# Pass previous screenshot path (different from state.latest_screenshot)
|
|
screenshot_path="/tmp/previous_screenshot.jpg",
|
|
state=mock_state,
|
|
minimal_list="",
|
|
)
|
|
|
|
# Ensure draw_dots was called with previous bounds [100,100][200,200] -> [150, 150]
|
|
# instead of current bounds [300,300][400,400] -> [350, 350]
|
|
mock_draw_dots.assert_called_once()
|
|
points_arg = mock_draw_dots.call_args[0][1]
|
|
assert points_arg == [[150, 150]] # Uses previous screenshot bounds!
|
|
assert points_arg != [[350, 350]] # Does not use current screenshot bounds!
|
|
|
|
labels_arg = mock_draw_dots.call_args[0][2]
|
|
assert labels_arg == ["1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_initial_visual_marking_previous_screenshot_no_ui_tree():
|
|
|
|
# Mock context and state
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/current_screenshot.jpg"
|
|
mock_state.latest_ui_hierarchy = [{"text": "Current Button", "bounds": "[300,300][400,400]"}]
|
|
|
|
# Mock Gemini Client and response
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.text = ""
|
|
|
|
mock_func_call = MagicMock()
|
|
mock_func_call.name = "submit_answer"
|
|
mock_func_call.args = {
|
|
"candidates": [{"label": "1", "coords": [500, 500], "description": "Visual target"}]
|
|
}
|
|
mock_response.function_calls = [mock_func_call]
|
|
mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
|
|
|
|
# Mock storage manager to return None (meaning not found in DB)
|
|
mock_storage = MagicMock()
|
|
mock_storage.get_image.return_value = None
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_raw_data",
|
|
b"",
|
|
b"fake_image_raw_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots,
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
# Pass previous screenshot path (different from state.latest_screenshot)
|
|
screenshot_path="/tmp/previous_screenshot.jpg",
|
|
state=mock_state,
|
|
minimal_list="",
|
|
)
|
|
|
|
# Ensure draw_dots was NEVER called since no ui_tree was available for previous screenshot
|
|
mock_draw_dots.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_initial_visual_marking_previous_screenshot_ocr_fusion():
|
|
from artemis.data_engine.models import ImageRecord
|
|
|
|
# Mock context and state
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/current_screenshot.jpg"
|
|
|
|
# Mock Gemini Client and response
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.text = ""
|
|
|
|
mock_func_call = MagicMock()
|
|
mock_func_call.name = "submit_answer"
|
|
mock_func_call.args = {
|
|
"candidates": [{"label": "1", "coords": [500, 500], "description": "Visual target"}]
|
|
}
|
|
mock_response.function_calls = [mock_func_call]
|
|
mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
|
|
|
|
# Mock storage manager
|
|
mock_storage = MagicMock()
|
|
|
|
# Mock ImageRecord with previous raw XML and raw OCR
|
|
mock_record = MagicMock(spec=ImageRecord)
|
|
mock_record.ui_tree = [{"text": "", "bounds": "[100,100][200,200]"}]
|
|
# First ocr_result is full screen text and ignored; second overlaps with bounds
|
|
mock_record.ocr_result = [
|
|
{"text": "full screen", "position": [{"x": 0, "y": 0}]},
|
|
{
|
|
"text": "OCR Text",
|
|
"position": [
|
|
{"x": 110, "y": 110},
|
|
{"x": 190, "y": 110},
|
|
{"x": 190, "y": 190},
|
|
{"x": 110, "y": 190},
|
|
],
|
|
},
|
|
]
|
|
mock_storage.get_image.return_value = mock_record
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_marked_data",
|
|
b"",
|
|
b"fake_image_marked_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots,
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
screenshot_path="/tmp/previous_screenshot.jpg",
|
|
state=mock_state,
|
|
minimal_list="",
|
|
)
|
|
|
|
# Ensure draw_dots was called with center coordinates of the fused OCR element [110,110][190,190] -> [150, 150]
|
|
mock_draw_dots.assert_called_once()
|
|
points_arg = mock_draw_dots.call_args[0][1]
|
|
assert points_arg == [[150, 150]] # Successfully fused and drawn!
|
|
|
|
labels_arg = mock_draw_dots.call_args[0][2]
|
|
assert labels_arg == ["1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_initial_visual_marking_previous_screenshot_on_the_fly_ocr():
|
|
from artemis.data_engine.models import ImageRecord
|
|
|
|
# Mock context and state
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/current_screenshot.jpg"
|
|
|
|
# Mock Gemini Client and response
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.text = ""
|
|
|
|
mock_func_call = MagicMock()
|
|
mock_func_call.name = "submit_answer"
|
|
mock_func_call.args = {
|
|
"candidates": [{"label": "1", "coords": [500, 500], "description": "Visual target"}]
|
|
}
|
|
mock_response.function_calls = [mock_func_call]
|
|
mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
|
|
|
|
# Mock storage manager
|
|
mock_storage = MagicMock()
|
|
|
|
# Mock ImageRecord containing ui_tree but ocr_result is None
|
|
mock_record = MagicMock(spec=ImageRecord)
|
|
mock_record.ui_tree = [{"text": "", "bounds": "[100,100][200,200]"}]
|
|
mock_record.ocr_result = None
|
|
mock_storage.get_image.return_value = mock_record
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
# Mock dynamic OCR result
|
|
mock_ocr_result = [
|
|
{"text": "full screen", "position": [{"x": 0, "y": 0}]},
|
|
{
|
|
"text": "On-the-fly OCR Text",
|
|
"position": [
|
|
{"x": 110, "y": 110},
|
|
{"x": 190, "y": 110},
|
|
{"x": 190, "y": 190},
|
|
{"x": 110, "y": 190},
|
|
],
|
|
},
|
|
]
|
|
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_marked_data",
|
|
b"",
|
|
b"fake_image_marked_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.perform_ocr",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_ocr_result,
|
|
) as mock_perform_ocr,
|
|
patch(
|
|
"artemis.agents.explorer.explorer.is_ocr_configured",
|
|
return_value=True,
|
|
),
|
|
patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots,
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
screenshot_path="/tmp/previous_screenshot.jpg",
|
|
state=mock_state,
|
|
minimal_list="",
|
|
)
|
|
|
|
# Verify perform_ocr was triggered on-the-fly
|
|
mock_perform_ocr.assert_called_once()
|
|
|
|
# Ensure draw_dots was called with center coordinates of the fused OCR element [110,110][190,190] -> [150, 150]
|
|
mock_draw_dots.assert_called_once()
|
|
points_arg = mock_draw_dots.call_args[0][1]
|
|
assert points_arg == [[150, 150]]
|
|
|
|
labels_arg = mock_draw_dots.call_args[0][2]
|
|
assert labels_arg == ["1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_denylisted_tool():
|
|
# Mock context and state
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
mock_ctx.agent_config = MagicMock()
|
|
mock_ctx.agent_config.denylisted_tools = {"explorer": ["search_xml_ocr"]}
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/test_screenshot.jpg"
|
|
|
|
# Response 1: Model calls search_xml_ocr (which is denylisted)
|
|
mock_response1 = MagicMock()
|
|
mock_response1.text = ""
|
|
mock_func_call1 = MagicMock()
|
|
mock_func_call1.name = "search_xml_ocr"
|
|
mock_func_call1.args = {"use_coords": False, "search_value": "test"}
|
|
mock_response1.function_calls = [mock_func_call1]
|
|
mock_response1.candidates = [MagicMock()]
|
|
mock_response1.candidates[0].content = types.Content(
|
|
role="model",
|
|
parts=[types.Part.from_function_call(name="search_xml_ocr", args=mock_func_call1.args)],
|
|
)
|
|
|
|
# Response 2: Model calls submit_answer
|
|
mock_response2 = MagicMock()
|
|
mock_response2.text = ""
|
|
mock_func_call2 = MagicMock()
|
|
mock_func_call2.name = "submit_answer"
|
|
mock_func_call2.args = {
|
|
"candidates": [
|
|
{
|
|
"label": "S1",
|
|
"coords": [500, 500],
|
|
"description": "Corrected button",
|
|
}
|
|
],
|
|
"fallback_message": "",
|
|
}
|
|
mock_response2.function_calls = [mock_func_call2]
|
|
mock_response2.candidates = [MagicMock()]
|
|
mock_response2.candidates[0].content = types.Content(
|
|
role="model",
|
|
parts=[types.Part.from_function_call(name="submit_answer", args=mock_func_call2.args)],
|
|
)
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.aio = MagicMock()
|
|
mock_client.aio.models = MagicMock()
|
|
mock_client.aio.models.generate_content = AsyncMock(
|
|
side_effect=[mock_response1, mock_response2]
|
|
)
|
|
|
|
# Mock storage manager
|
|
mock_storage = MagicMock()
|
|
mock_storage.get_image.return_value = "dummy_hash"
|
|
|
|
# Mock file reading for prompt
|
|
mock_prompt = MOCK_PROMPT_JSON
|
|
|
|
def custom_open(file, mode="r", *args, **kwargs):
|
|
mock_file = MagicMock()
|
|
if "b" in mode:
|
|
mock_file.read.side_effect = [
|
|
b"fake_image_data",
|
|
b"",
|
|
b"fake_image_data",
|
|
]
|
|
else:
|
|
mock_file.read.return_value = '{"templates": [], "instructions": ""}'
|
|
|
|
mock_file.__enter__.return_value = mock_file
|
|
return mock_file
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.genai.Client",
|
|
return_value=mock_client,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value=mock_prompt),
|
|
patch("builtins.open", custom_open),
|
|
):
|
|
explorer = Explorer(mock_ctx)
|
|
|
|
# Verify get_exposed_tools doesn't expose search_xml_ocr
|
|
exposed_tools = explorer.get_exposed_tools()
|
|
exposed_tool_names = [tool.name for tool in exposed_tools]
|
|
assert "search_xml_ocr" not in exposed_tool_names
|
|
|
|
result = await explorer.run(
|
|
query="Find button",
|
|
context_feedback="",
|
|
screenshot_path="/tmp/test_screenshot.jpg",
|
|
minimal_list="",
|
|
state=mock_state,
|
|
)
|
|
|
|
expected_outcome = json.dumps(mock_func_call2.args, ensure_ascii=False)
|
|
assert result == expected_outcome
|
|
|
|
# Ensure we called generate_content twice
|
|
assert mock_client.aio.models.generate_content.call_count == 2
|
|
|
|
# Verify the contents passed to the second call contain the error response
|
|
# indicating that the tool is denylisted and unavailable.
|
|
call_args = mock_client.aio.models.generate_content.call_args_list[1]
|
|
contents = call_args[1]["contents"]
|
|
|
|
# Find the tool response in contents
|
|
tool_content = next(
|
|
(
|
|
c
|
|
for c in contents
|
|
if c.role in ("tool", "user") and any(p.function_response for p in c.parts)
|
|
),
|
|
None,
|
|
)
|
|
assert tool_content is not None
|
|
|
|
# Verify that the response part contains the expected error message
|
|
found_error = False
|
|
for part in tool_content.parts:
|
|
func_resp = part.function_response
|
|
if func_resp and func_resp.name == "search_xml_ocr":
|
|
assert "denylisted and unavailable" in func_resp.response["error"]
|
|
found_error = True
|
|
|
|
assert found_error
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_inspect_region():
|
|
import pathlib
|
|
import numpy as np
|
|
|
|
# Mock context
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
|
|
explorer = Explorer(mock_ctx)
|
|
explorer.width = 1080
|
|
explorer.height = 2400
|
|
explorer.screenshot_path = "/tmp/test_screenshot.jpg"
|
|
explorer.image_name = "test_image"
|
|
|
|
dummy_img = np.zeros((2400, 1080, 3), dtype=np.uint8)
|
|
|
|
with (
|
|
patch("cv2.imread", return_value=dummy_img) as mock_imread,
|
|
patch("cv2.resize", return_value=dummy_img) as mock_resize,
|
|
patch("cv2.imwrite", return_value=True) as mock_imwrite,
|
|
patch("pathlib.Path.mkdir") as mock_mkdir,
|
|
patch("glob.glob", return_value=[]),
|
|
patch("artemis.agents.explorer.explorer.settings") as mock_settings,
|
|
):
|
|
mock_settings.TRACES_PATH = pathlib.Path("/tmp/traces")
|
|
res = await explorer.exec_inspect_region(
|
|
x_min=100, y_min=200, x_max=300, y_max=400, zoom_factor=2.0
|
|
)
|
|
assert res["image_path"] is not None
|
|
assert "Inspected region coordinates" in res["text"]
|
|
mock_imread.assert_called_once_with("/tmp/test_screenshot.jpg")
|
|
|
|
# Verify resize calculations:
|
|
# px_start = 100 * 1080 / 1000 = 108
|
|
# py_start = 200 * 2400 / 1000 = 480
|
|
# px_end = 300 * 1080 / 1000 = 324
|
|
# py_end = 400 * 2400 / 1000 = 960
|
|
# width = 324 - 108 = 216
|
|
# height = 960 - 480 = 480
|
|
# new_width = 216 * 2.0 = 432
|
|
# new_height = 480 * 2.0 = 960
|
|
mock_resize.assert_called_once()
|
|
resize_args = mock_resize.call_args[0]
|
|
assert resize_args[1] == (432, 960)
|
|
mock_imwrite.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_final_turn_tool_stripping():
|
|
"""Test that on the final iteration (turn 10), Explorer strips other tools and injects warning."""
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.device = MagicMock()
|
|
mock_ctx.device.device_width = 1080
|
|
mock_ctx.device.device_height = 2400
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
explorer = Explorer(mock_ctx)
|
|
|
|
mock_ctx.llm_config = MagicMock()
|
|
mock_ctx.llm_config.explorer = MagicMock()
|
|
mock_ctx.llm_config.explorer.model = "gemini-3.7-flash"
|
|
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.latest_screenshot = "/tmp/test_screenshot.jpg"
|
|
|
|
# Set up generate_content mock to capture call arguments
|
|
mock_client = MagicMock()
|
|
mock_ctx._genai_client = mock_client
|
|
|
|
# Turn 1: model returns search_xml_ocr
|
|
fc_search = MagicMock()
|
|
fc_search.name = "search_xml_ocr"
|
|
fc_search.args = {"use_coords": False, "search_value": "test"}
|
|
mock_response = MagicMock()
|
|
mock_response.function_calls = [fc_search]
|
|
mock_response.candidates = [MagicMock()]
|
|
mock_response.candidates[0].content = types.Content(
|
|
role="model",
|
|
parts=[types.Part.from_function_call(name="search_xml_ocr", args=fc_search.args)],
|
|
)
|
|
|
|
# Always return the same mock response to force it to run out of turns
|
|
mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
|
|
|
|
from unittest.mock import mock_open
|
|
|
|
with (
|
|
patch("builtins.open", mock_open(read_data=b"")),
|
|
patch("hashlib.sha256") as mock_sha,
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_text", return_value="{}"),
|
|
):
|
|
mock_sha.return_value.hexdigest.return_value = "dummy_hash"
|
|
|
|
await explorer.run(
|
|
query="test",
|
|
context_feedback="",
|
|
screenshot_path="/tmp/test_screenshot.jpg",
|
|
state=mock_state,
|
|
minimal_list="test list",
|
|
)
|
|
|
|
# Total number of iterations is max_iterations = 3 under "pro" mode default
|
|
assert mock_client.aio.models.generate_content.call_count == 3
|
|
|
|
# Check the final turn call (index 2) keyword args
|
|
call_args_3 = mock_client.aio.models.generate_content.call_args_list[2]
|
|
config = call_args_3.kwargs["config"]
|
|
|
|
# Tools in the 10th call should only contain "submit_answer"
|
|
tools = config.tools[0].function_declarations
|
|
assert len(tools) == 1
|
|
assert tools[0].name == "submit_answer"
|
|
|
|
# Warning should be injected in contents
|
|
contents = call_args_3.kwargs["contents"]
|
|
# Find warning in contents parts
|
|
warning_found = False
|
|
for content in contents:
|
|
if content.role == "user":
|
|
for part in content.parts:
|
|
if part.text and "[WARNING] This is your final iteration" in part.text:
|
|
warning_found = True
|
|
assert warning_found, "Warning message was not found in contents"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_explorer_ask_perception_tool():
|
|
import numpy as np
|
|
from unittest.mock import mock_open
|
|
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = "/tmp/test_session"
|
|
explorer = Explorer(mock_ctx)
|
|
explorer.image_name = "dummy_hash"
|
|
explorer.screenshot_path = "/tmp/test_screenshot.jpg"
|
|
explorer.width = 1080
|
|
explorer.height = 2400
|
|
explorer.global_label_idx = 1
|
|
explorer.image_pool = {
|
|
"img_0": {
|
|
"path": "/tmp/test_screenshot.jpg",
|
|
"transform": {
|
|
"offset_x": 0.0,
|
|
"offset_y": 0.0,
|
|
"scale_x": 1.0,
|
|
"scale_y": 1.0,
|
|
},
|
|
"description": "Original complete screenshot",
|
|
}
|
|
}
|
|
|
|
# Mock search_ui_func
|
|
mock_search_ui_res = {
|
|
"matches": [
|
|
{
|
|
"matched_text": "Settings",
|
|
"bounds": [540, 480, 540, 480],
|
|
"type": "xml",
|
|
},
|
|
{
|
|
"matched_text": "Sign In",
|
|
"bounds": [100, 100, 200, 200],
|
|
"type": "ocr",
|
|
},
|
|
]
|
|
}
|
|
|
|
# Mock storage record
|
|
mock_record = MagicMock()
|
|
mock_record.ui_tree = [
|
|
{
|
|
"bounds": "[500,400][600,560]",
|
|
"text": "Settings Button",
|
|
"resource-id": "com.android.settings:id/btn_settings",
|
|
"class": "android.widget.Button",
|
|
}
|
|
]
|
|
mock_record.ocr_result = []
|
|
|
|
mock_storage = MagicMock()
|
|
mock_storage.get_image.return_value = mock_record
|
|
|
|
# Mock _run_object_detection
|
|
mock_detector_res = {
|
|
"detected": [{"label": "profile icon", "point": [900, 100]}],
|
|
"failed": [],
|
|
}
|
|
|
|
# Mock draw_dots
|
|
mock_draw_dots = MagicMock()
|
|
|
|
with (
|
|
patch(
|
|
"artemis.agents.explorer.explorer.search_ui_func",
|
|
return_value=mock_search_ui_res,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch(
|
|
"artemis.mcp.xml_search_server.StorageManager",
|
|
return_value=mock_storage,
|
|
),
|
|
patch(
|
|
"artemis.agents.explorer.explorer._run_object_detection",
|
|
return_value=mock_detector_res,
|
|
),
|
|
patch("artemis.agents.explorer.explorer.draw_dots", mock_draw_dots),
|
|
patch("artemis.utils.visualization.draw_dots", mock_draw_dots),
|
|
patch("pathlib.Path.mkdir"),
|
|
patch("glob.glob", return_value=[]),
|
|
patch("builtins.open", mock_open(read_data=b'{"templates": []}')),
|
|
patch("cv2.imread", return_value=np.zeros((2400, 1080, 3), dtype=np.uint8)),
|
|
):
|
|
result = await explorer.exec_ask_perception_tool(
|
|
search_query="Settings",
|
|
nx=550,
|
|
ny=200,
|
|
detect_queries=["profile icon"],
|
|
)
|
|
|
|
assert "Text Search Results are:" in result["text"]
|
|
assert "Coordinate Search Results are:" in result["text"]
|
|
assert "Object Detection Results are:" in result["text"]
|
|
|
|
# XML Text search should return label X1
|
|
assert "[X1]" in result["text"]
|
|
# OCR Text search should return label O2
|
|
assert "[O2]" in result["text"]
|
|
# Coordinate search should return label X3
|
|
assert "[X3]" in result["text"]
|
|
# Object detection should return label D4 (since global label idx is shared/contiguous)
|
|
assert "[D4]" in result["text"]
|
|
|
|
assert len(result["image_paths"]) == 3
|
|
|
|
# Verify call arguments of draw_dots
|
|
assert mock_draw_dots.call_count == 3
|
|
# First call (XML/OCR Text search): green color
|
|
first_call = mock_draw_dots.call_args_list[0]
|
|
assert first_call.kwargs.get("color") == "green"
|
|
assert first_call.args[2] == ["X1", "O2"]
|
|
|
|
# Second call (XML Coordinate audit): blue color
|
|
second_call = mock_draw_dots.call_args_list[1]
|
|
assert second_call.kwargs.get("color") == "blue"
|
|
assert second_call.args[2] == ["X3"]
|
|
|
|
# Third call (Object detection): red color
|
|
third_call = mock_draw_dots.call_args_list[2]
|
|
assert third_call.kwargs.get("color") == "red"
|
|
assert third_call.args[2] == ["D4"]
|
|
|
|
|
|
def test_explorer_prune_historical_images():
|
|
from artemis.agents.explorer.explorer import Explorer
|
|
from google.genai import types
|
|
|
|
mock_ctx = MagicMock()
|
|
explorer = Explorer(mock_ctx)
|
|
|
|
part1 = types.Part(inline_data=types.Blob(data=b"img1", mime_type="image/jpeg"))
|
|
part2 = types.Part(inline_data=types.Blob(data=b"img2", mime_type="image/jpeg"))
|
|
part3 = types.Part.from_text(text="Some text")
|
|
|
|
contents = [
|
|
types.Content(role="user", parts=[part1, part3]),
|
|
types.Content(role="user", parts=[part2]),
|
|
]
|
|
|
|
explorer._prune_historical_images(contents, keep_last=1)
|
|
|
|
# First image part should be pruned to text part
|
|
assert contents[0].parts[0].text == "[Image pruned to maintain visual focus on latest state]"
|
|
assert contents[0].parts[1].text == "Some text"
|
|
# Second (last) image part should be kept intact
|
|
assert contents[1].parts[0] is part2
|
|
|
|
|
|
class _FakeGenaiError(Exception):
|
|
"""Mimics google.genai.errors.APIError's shape (a ``code`` attribute)."""
|
|
|
|
def __init__(self, message: str, code: int | None = None):
|
|
super().__init__(message)
|
|
self.code = code
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_reliability_permanent_error_is_not_retried():
|
|
from artemis.agents.explorer.explorer import _generate_content_with_reliability
|
|
from artemis.llm.reliability import LLMPermanentError
|
|
|
|
calls = 0
|
|
|
|
async def op():
|
|
nonlocal calls
|
|
calls += 1
|
|
raise _FakeGenaiError("PERMISSION_DENIED: forbidden", code=403)
|
|
|
|
with pytest.raises(LLMPermanentError) as exc_info:
|
|
await _generate_content_with_reliability(op)
|
|
assert calls == 1
|
|
assert exc_info.value.failure.category.value == "authentication"
|
|
assert isinstance(exc_info.value.__cause__, _FakeGenaiError)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_reliability_bad_request_is_not_retried():
|
|
from artemis.agents.explorer.explorer import _generate_content_with_reliability
|
|
from artemis.llm.reliability import LLMPermanentError
|
|
|
|
calls = 0
|
|
|
|
async def op():
|
|
nonlocal calls
|
|
calls += 1
|
|
raise _FakeGenaiError("invalid request payload", code=400)
|
|
|
|
with pytest.raises(LLMPermanentError) as exc_info:
|
|
await _generate_content_with_reliability(op)
|
|
assert calls == 1
|
|
assert exc_info.value.failure.category.value == "bad_request"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_reliability_retries_transient_then_succeeds():
|
|
from artemis.agents.explorer import explorer as explorer_module
|
|
|
|
calls = 0
|
|
sentinel = object()
|
|
|
|
async def op():
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls < 2:
|
|
raise _FakeGenaiError("429 rate limit exceeded", code=429)
|
|
return sentinel
|
|
|
|
with patch.object(explorer_module.asyncio, "sleep", new=AsyncMock()) as sleep_mock:
|
|
result = await explorer_module._generate_content_with_reliability(op)
|
|
assert result is sentinel
|
|
assert calls == 2
|
|
assert sleep_mock.await_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_reliability_exhausts_to_typed_error():
|
|
from artemis.agents.explorer import explorer as explorer_module
|
|
from artemis.llm.reliability import (
|
|
FailureCategory,
|
|
LLMExhaustedError,
|
|
retry_policy_for,
|
|
)
|
|
|
|
calls = 0
|
|
|
|
async def op():
|
|
nonlocal calls
|
|
calls += 1
|
|
raise _FakeGenaiError("503 service unavailable", code=503)
|
|
|
|
with patch.object(explorer_module.asyncio, "sleep", new=AsyncMock()):
|
|
with pytest.raises(LLMExhaustedError) as exc_info:
|
|
await explorer_module._generate_content_with_reliability(op)
|
|
assert calls == retry_policy_for(FailureCategory.PROVIDER_UNAVAILABLE).max_attempts
|
|
assert exc_info.value.failure.category.value == "provider_unavailable"
|
|
assert isinstance(exc_info.value.__cause__, _FakeGenaiError)
|