chenmingxuan
2ddec53954
Localize the web console and persist custom model and capability settings. Use OpenAI Responses exclusively and move runtime limits into environment configuration. Bootstrap uv and the virtual environment on Windows, support optional China mirrors, and package source distributions without environments, build output, or credentials. Update regression tests and repository quality checks for the new behavior.
225 lines
7.2 KiB
Python
225 lines
7.2 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 base64
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from artemis.context import ArtemisContext
|
|
from artemis.drivers.base import BaseDeviceDriver
|
|
from artemis.graph.state import State
|
|
from artemis.tools.base import ArtemisTool
|
|
from artemis.tools.mobile.ocr import (
|
|
OCRRecognition,
|
|
OcrArgs,
|
|
OcrRecognition,
|
|
OcrRecognitionTool,
|
|
OcrTool,
|
|
get_ocr_tool,
|
|
ocr_recognition,
|
|
ocr_recognition_wrapper,
|
|
)
|
|
from langchain_core.messages import ToolMessage
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_ctx():
|
|
ctx = MagicMock(spec=ArtemisContext)
|
|
ctx.device = MagicMock()
|
|
ctx.device.device_width = 1080
|
|
ctx.device.device_height = 2400
|
|
return ctx
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_driver():
|
|
driver = MagicMock(spec=BaseDeviceDriver)
|
|
driver.screen_size = (1080, 2400)
|
|
screen_data = MagicMock()
|
|
screen_data.screenshot_base64 = base64.b64encode(b"dummy_image_data").decode("utf-8")
|
|
screen_data.screenshot_bytes = b"dummy_image_data"
|
|
driver.get_screen_data = AsyncMock(return_value=screen_data)
|
|
return driver
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_state(tmp_path):
|
|
screenshot_file = tmp_path / "test_screenshot.png"
|
|
screenshot_file.write_bytes(b"dummy_image_content")
|
|
|
|
state = MagicMock(spec=State)
|
|
state.latest_screenshot = str(screenshot_file)
|
|
return state
|
|
|
|
|
|
def test_ocr_tool_subclass():
|
|
"""Verify OcrRecognitionTool is a subclass of ArtemisTool."""
|
|
assert issubclass(OcrRecognitionTool, ArtemisTool)
|
|
assert issubclass(OcrRecognition, ArtemisTool)
|
|
assert issubclass(OCRRecognition, ArtemisTool)
|
|
assert issubclass(OcrTool, ArtemisTool)
|
|
assert isinstance(ocr_recognition, ArtemisTool)
|
|
assert isinstance(ocr_recognition, OcrRecognitionTool)
|
|
|
|
assert ocr_recognition.name == "ocr_recognition"
|
|
assert ocr_recognition.category == "explorer"
|
|
assert ocr_recognition.args_schema == OcrArgs
|
|
|
|
# GenAI FunctionDeclaration export
|
|
declaration = ocr_recognition.to_tool_declaration()
|
|
assert declaration.name == "ocr_recognition"
|
|
|
|
# Wrapper check
|
|
assert ocr_recognition_wrapper is not None
|
|
assert ocr_recognition_wrapper.tool_fn_getter == get_ocr_tool
|
|
|
|
|
|
class SimpleState:
|
|
def __init__(self, screenshot_path: str):
|
|
self.latest_screenshot = screenshot_path
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ocr_direct_execution_with_state(mock_ctx, tmp_path):
|
|
"""Verify direct execution with state and screenshot file."""
|
|
screenshot_file = tmp_path / "test_screenshot.png"
|
|
screenshot_file.write_bytes(b"dummy_image_content")
|
|
simple_state = SimpleState(str(screenshot_file))
|
|
|
|
mock_results = [
|
|
{
|
|
"text": "Submit",
|
|
"position": [
|
|
{"x": 100, "y": 200},
|
|
{"x": 200, "y": 200},
|
|
{"x": 200, "y": 300},
|
|
{"x": 100, "y": 300},
|
|
],
|
|
}
|
|
]
|
|
with patch(
|
|
"artemis.tools.mobile.ocr.perform_ocr",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_results,
|
|
):
|
|
result = await ocr_recognition.execute(ctx=mock_ctx, state=simple_state)
|
|
assert isinstance(result, ToolMessage)
|
|
assert "OCR Recognition successful." in result.content
|
|
assert "Submit" in result.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ocr_direct_execution_with_driver(mock_ctx, mock_driver):
|
|
"""Verify direct execution with driver."""
|
|
mock_results = [
|
|
{
|
|
"text": "Settings",
|
|
"position": [
|
|
{"x": 500, "y": 500},
|
|
{"x": 600, "y": 600},
|
|
],
|
|
}
|
|
]
|
|
with patch(
|
|
"artemis.tools.mobile.ocr.perform_ocr",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_results,
|
|
):
|
|
result = await ocr_recognition.execute(ctx=mock_ctx, driver=mock_driver)
|
|
assert "OCR Recognition successful." in result
|
|
assert "Settings" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ocr_no_text_detected(mock_ctx, tmp_path):
|
|
"""Verify outcome when no text is detected."""
|
|
screenshot_file = tmp_path / "test_empty_screenshot.png"
|
|
screenshot_file.write_bytes(b"dummy_empty_content")
|
|
simple_state = SimpleState(str(screenshot_file))
|
|
|
|
with patch(
|
|
"artemis.tools.mobile.ocr.perform_ocr",
|
|
new_callable=AsyncMock,
|
|
return_value=[],
|
|
):
|
|
result = await ocr_recognition.execute(ctx=mock_ctx, state=simple_state)
|
|
assert isinstance(result, ToolMessage)
|
|
assert "No text detected on the screen." in result.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ocr_with_state_command(mock_ctx, mock_state):
|
|
"""Verify OcrRecognitionTool returns ToolMessage when state is provided."""
|
|
mock_results = [{"text": "Hello", "position": [{"x": 100, "y": 100}]}]
|
|
with patch(
|
|
"artemis.tools.mobile.ocr.perform_ocr",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_results,
|
|
):
|
|
cmd = await ocr_recognition.execute(
|
|
ctx=mock_ctx,
|
|
tool_call_id="call_ocr_1",
|
|
state=mock_state,
|
|
)
|
|
assert isinstance(cmd, ToolMessage)
|
|
assert cmd.tool_call_id == "call_ocr_1"
|
|
assert cmd.status == "success"
|
|
assert "Hello" in cmd.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ocr_failure_missing_screenshot(mock_ctx):
|
|
"""Verify error handling when screenshot path is missing."""
|
|
empty_state = MagicMock(spec=State)
|
|
empty_state.latest_screenshot = None
|
|
|
|
cmd = await ocr_recognition.execute(
|
|
ctx=mock_ctx,
|
|
tool_call_id="call_ocr_err",
|
|
state=empty_state,
|
|
)
|
|
assert isinstance(cmd, ToolMessage)
|
|
assert cmd.status == "error"
|
|
assert "OCR Recognition failed" in cmd.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ocr_callable_execution(mock_ctx, mock_driver):
|
|
"""Verify invoking ocr_recognition directly as a callable."""
|
|
with patch(
|
|
"artemis.tools.mobile.ocr.perform_ocr",
|
|
new_callable=AsyncMock,
|
|
return_value=[{"text": "Callable", "position": []}],
|
|
):
|
|
result = await ocr_recognition(ctx=mock_ctx, driver=mock_driver)
|
|
assert "OCR Recognition successful." in result
|
|
assert "Callable" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_ocr_tool_langchain_ainvoke(mock_ctx, mock_driver):
|
|
"""Verify get_ocr_tool exports a LangChain tool that works with ainvoke."""
|
|
mock_ctx._active_driver = mock_driver
|
|
ocr_tool = get_ocr_tool(mock_ctx)
|
|
assert ocr_tool.name == "ocr_recognition"
|
|
|
|
with patch(
|
|
"artemis.tools.mobile.ocr.perform_ocr",
|
|
new_callable=AsyncMock,
|
|
return_value=[{"text": "LangChain", "position": []}],
|
|
):
|
|
result = await ocr_tool.ainvoke({})
|
|
assert "OCR Recognition successful." in result
|
|
assert "LangChain" in result
|