308 lines
12 KiB
Python
308 lines
12 KiB
Python
"""Android artifact integrity, path separation and Flash/Pro tool parity."""
|
|
|
|
import asyncio
|
|
import hashlib
|
|
from io import BytesIO
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
from PIL import Image
|
|
import pytest
|
|
|
|
from artemis.agents.flash.runner import FlashRunner
|
|
from artemis.agents.validator.tool_declarations import ToolExecutionResult
|
|
from artemis.core.tool_failure import is_tool_failure
|
|
from artemis.mcp.action_executor import McpActionExecutor
|
|
from artemis.tools import android_files as files
|
|
|
|
|
|
@pytest.fixture
|
|
def ctx(tmp_path):
|
|
return SimpleNamespace(
|
|
execution_target="mixed",
|
|
device=SimpleNamespace(device_id="task-device"),
|
|
data_engine=SimpleNamespace(base_dir=tmp_path / "trace", current_session_id="session"),
|
|
execution_setup=None,
|
|
actuator=None,
|
|
plugins=None,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def png():
|
|
stream = BytesIO()
|
|
Image.new("RGB", (144, 312), "red").save(stream, "PNG")
|
|
return stream.getvalue()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_screenshot_saved_native_and_desktop_copy_verified(ctx, png, tmp_path, monkeypatch):
|
|
desktop = tmp_path / "重定向桌面"
|
|
desktop.mkdir()
|
|
monkeypatch.setattr(files, "desktop_directory", lambda: desktop)
|
|
run = AsyncMock(return_value=png)
|
|
monkeypatch.setattr(files.AndroidFiles, "_run", run)
|
|
result = json.loads(
|
|
await files.android_file.execute(
|
|
ctx=ctx, action="screenshot", filename="手机截图.png", destination="desktop"
|
|
)
|
|
)
|
|
assert result["ok"] and result["verified"]
|
|
assert result["width"] == 144 and result["height"] == 312
|
|
assert result["device_serial"] == "task-device"
|
|
assert Path(result["path"]) == desktop / "手机截图.png"
|
|
assert Path(result["path"]).read_bytes() == Path(result["trace_path"]).read_bytes() == png
|
|
assert result["sha256"] == hashlib.sha256(png).hexdigest()
|
|
run.assert_awaited_once_with("exec-out", "screencap", "-p", limit=files.MAX_FILE_BYTES)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_existing_artifact_is_preserved_unless_explicit_overwrite(ctx, png, monkeypatch):
|
|
monkeypatch.setattr(files.AndroidFiles, "_run", AsyncMock(return_value=png))
|
|
first = json.loads(
|
|
await files.android_file.execute(ctx=ctx, action="screenshot", filename="a.png")
|
|
)
|
|
path = Path(first["trace_path"])
|
|
path.write_bytes(b"original")
|
|
error = await files.android_file.execute(ctx=ctx, action="screenshot", filename="a.png")
|
|
assert is_tool_failure(error) and path.read_bytes() == b"original"
|
|
result = await files.android_file.execute(
|
|
ctx=ctx, action="screenshot", filename="a.png", overwrite=True
|
|
)
|
|
assert json.loads(result)["ok"] and path.read_bytes() == png
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"name", ["../escape.png", "C:\\out.png", "x:y.png", "CON.png", "x.png ", "a\n.png"]
|
|
)
|
|
async def test_invalid_host_paths_fail_before_adb(ctx, name, monkeypatch):
|
|
run = AsyncMock()
|
|
monkeypatch.setattr(files.AndroidFiles, "_run", run)
|
|
result = await files.android_file.execute(ctx=ctx, action="screenshot", filename=name)
|
|
assert is_tool_failure(result)
|
|
run.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"device_path", ["C:\\Desktop\\x", "relative", "/sdcard/../x", "/", "/sdcard/a\n"]
|
|
)
|
|
async def test_android_path_validation(ctx, device_path, monkeypatch):
|
|
run = AsyncMock()
|
|
monkeypatch.setattr(files.AndroidFiles, "_run", run)
|
|
assert is_tool_failure(
|
|
await files.android_file.execute(ctx=ctx, action="read_text", device_path=device_path)
|
|
)
|
|
run.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalid_image_does_not_create_artifact(ctx, monkeypatch):
|
|
monkeypatch.setattr(files.AndroidFiles, "_run", AsyncMock(return_value=b"screencap failed"))
|
|
assert is_tool_failure(await files.android_file.execute(ctx=ctx, action="screenshot"))
|
|
assert not (ctx.data_engine.base_dir / "artifacts").exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_and_pull_verify_chinese_and_binary_data(ctx, monkeypatch):
|
|
service = files.AndroidFiles(ctx)
|
|
data = "中文内容\n第二行".encode()
|
|
digest = hashlib.sha256(data).hexdigest()
|
|
shell = AsyncMock(side_effect=[str(len(data)), digest, digest])
|
|
monkeypatch.setattr(service, "_shell", shell)
|
|
monkeypatch.setattr(service, "_run", AsyncMock(return_value=data))
|
|
result = await service.execute(
|
|
files.AndroidFileArgs(action="read_text", device_path="/sdcard/a '中'.txt")
|
|
)
|
|
assert result["content"] == data.decode()
|
|
assert "'\"'\"'" in shell.call_args_list[0].args[0] # POSIX shell quoting
|
|
binary = b"\x00\xff\r\n\x1a"
|
|
digest = hashlib.sha256(binary).hexdigest()
|
|
shell.side_effect = [str(len(binary)), digest, digest]
|
|
service._run.return_value = binary
|
|
result = await service.execute(
|
|
files.AndroidFileArgs(action="pull", device_path="/sdcard/file.bin")
|
|
)
|
|
assert Path(result["path"]).read_bytes() == binary
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_changed_or_oversize_file_is_not_saved(ctx, monkeypatch):
|
|
service = files.AndroidFiles(ctx)
|
|
monkeypatch.setattr(service, "_shell", AsyncMock(side_effect=["3", "0" * 64]))
|
|
run = AsyncMock(return_value=b"abc")
|
|
monkeypatch.setattr(service, "_run", run)
|
|
with pytest.raises(OSError, match="传输"):
|
|
await service.execute(files.AndroidFileArgs(action="pull", device_path="/sdcard/a"))
|
|
assert not service.root.exists()
|
|
service._shell.side_effect = [str(files.MAX_TEXT_BYTES + 1)]
|
|
run.reset_mock()
|
|
with pytest.raises(ValueError, match="上限"):
|
|
await service.execute(files.AndroidFileArgs(action="read_text", device_path="/sdcard/a"))
|
|
run.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_write_text_pushes_utf8_and_verifies_before_rename(ctx, monkeypatch):
|
|
service = files.AndroidFiles(ctx)
|
|
data = "第一行\n第二行".encode()
|
|
calls = []
|
|
|
|
async def run(*args, **kwargs):
|
|
if args[0] == "push":
|
|
assert Path(args[1]).read_bytes() == data
|
|
calls.append("push")
|
|
return b""
|
|
|
|
async def shell(command):
|
|
calls.append(command)
|
|
return hashlib.sha256(data).hexdigest() if command.startswith("sha256sum") else ""
|
|
|
|
monkeypatch.setattr(service, "_run", run)
|
|
monkeypatch.setattr(service, "_shell", shell)
|
|
result = await service.execute(
|
|
files.AndroidFileArgs(
|
|
action="write_text", device_path="/sdcard/Download/文本.txt", content=data.decode()
|
|
)
|
|
)
|
|
assert result["verified"] and result["bytes"] == len(data)
|
|
assert any("mv -n" in call and "test ! -e" in call for call in calls)
|
|
assert calls.index("push") < next(i for i, c in enumerate(calls) if "mv -n" in c)
|
|
assert not list(service.root.iterdir()) # no leaked local staging files
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_permission_failure_propagates_and_does_not_claim_success(ctx, monkeypatch):
|
|
monkeypatch.setattr(
|
|
files.AndroidFiles, "_run", AsyncMock(side_effect=OSError("Permission denied"))
|
|
)
|
|
result = await files.android_file.execute(
|
|
ctx=ctx, action="read_text", device_path="/data/private.txt"
|
|
)
|
|
assert is_tool_failure(result) and "Permission denied" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_existing_android_file_is_never_pushed_over(ctx, monkeypatch):
|
|
service = files.AndroidFiles(ctx)
|
|
monkeypatch.setattr(service, "_shell", AsyncMock(return_value="exists"))
|
|
push = AsyncMock()
|
|
monkeypatch.setattr(service, "_run", push)
|
|
with pytest.raises(FileExistsError, match="设备文件已存在"):
|
|
await service.execute(
|
|
files.AndroidFileArgs(action="write_text", device_path="/sdcard/a", content="new")
|
|
)
|
|
push.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure", ["cancel", "timeout", "limit"])
|
|
async def test_adb_subprocess_is_reaped_on_interrupt(ctx, monkeypatch, failure):
|
|
service = files.AndroidFiles(ctx)
|
|
process = Mock(returncode=None)
|
|
process.stdout.read = AsyncMock()
|
|
process.stderr.read = AsyncMock(return_value=b"")
|
|
if failure == "limit":
|
|
process.stdout.read.return_value = b"too large"
|
|
expected = ValueError
|
|
else:
|
|
|
|
async def wait_forever(*args):
|
|
await asyncio.Event().wait()
|
|
|
|
process.stdout.read.side_effect = wait_forever
|
|
expected = asyncio.CancelledError if failure == "cancel" else TimeoutError
|
|
process.wait = AsyncMock(return_value=0)
|
|
|
|
def kill():
|
|
process.returncode = -1
|
|
|
|
process.kill.side_effect = kill
|
|
spawn = AsyncMock(return_value=process)
|
|
monkeypatch.setattr(asyncio, "create_subprocess_exec", spawn)
|
|
monkeypatch.setattr(files, "ADB_TIMEOUT_SECONDS", 0.02)
|
|
task = asyncio.create_task(service._run("exec-out", "cat /sdcard/a", limit=2))
|
|
await asyncio.sleep(0)
|
|
if failure == "cancel":
|
|
task.cancel()
|
|
with pytest.raises(expected):
|
|
await task
|
|
process.kill.assert_called_once()
|
|
process.wait.assert_awaited()
|
|
argv = spawn.call_args.args
|
|
assert argv[argv.index("-s") + 1] == "task-device"
|
|
assert "-H" in argv and "-P" in argv
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flash_and_pro_share_schema_handler_and_failure_status(ctx, png, monkeypatch):
|
|
monkeypatch.setattr(files.AndroidFiles, "_run", AsyncMock(return_value=png))
|
|
monkeypatch.setattr("artemis.tools.base.get_driver", lambda ctx: None)
|
|
monkeypatch.setattr(
|
|
"artemis.mcp.action_executor.TraceSpan",
|
|
lambda **kw: Mock(__enter__=lambda self: self, __exit__=lambda *a: False),
|
|
)
|
|
runner = object.__new__(FlashRunner)
|
|
runner.ctx = ctx
|
|
assert "android_file" in {t.name for t in runner._get_tools()}
|
|
pro = files.android_file.to_langchain_tool(ctx)
|
|
assert pro.args_schema is files.AndroidFileArgs
|
|
result = await pro.ainvoke({"action": "screenshot", "filename": "pro.png"})
|
|
assert json.loads(result)["verified"]
|
|
executor = McpActionExecutor(ctx, actuator=Mock())
|
|
result = await executor.execute(
|
|
"android_file", {"action": "screenshot", "filename": "flash.png"}, "call", None
|
|
)
|
|
assert result.status == "success" and json.loads(result.text_summary)["verified"]
|
|
result = await executor.execute(
|
|
"android_file", {"action": "screenshot", "filename": "flash.png"}, "call2", None
|
|
)
|
|
assert result.status == "error"
|
|
ctx.execution_target = "computer"
|
|
assert "android_file" not in {t.name for t in runner._get_tools()}
|
|
assert is_tool_failure(await files.android_file.execute(ctx=ctx, action="screenshot"))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("status", ["success", "error"])
|
|
async def test_flash_records_file_receipts_as_steps_for_final_verification(ctx, status):
|
|
runner = object.__new__(FlashRunner)
|
|
runner.ctx = ctx
|
|
runner._llm = None
|
|
runner.summarizer = None
|
|
result = ToolExecutionResult(
|
|
tool_call_id="file-call",
|
|
tool_name="android_file",
|
|
status=status,
|
|
text_summary='{"verified": true, "sha256": "receipt"}'
|
|
if status == "success"
|
|
else "Permission denied",
|
|
)
|
|
runner.executor = SimpleNamespace(
|
|
action_tool_names=frozenset({"click"}), execute=AsyncMock(return_value=result)
|
|
)
|
|
runner._record_action_step = Mock(return_value="recorded-file-step")
|
|
turn = SimpleNamespace(step_keys=[], actions=[])
|
|
messages = []
|
|
await runner._execute_and_record_action(
|
|
name="android_file",
|
|
args={"action": "screenshot"},
|
|
tc_id="file-call",
|
|
state=SimpleNamespace(),
|
|
messages=messages,
|
|
raw_text="保存当前安卓屏幕",
|
|
step_token_usage={},
|
|
pre_screenshot_bytes=b"previous-observation",
|
|
xml_list="UI",
|
|
action_sequence=0,
|
|
turn=turn,
|
|
)
|
|
runner._record_action_step.assert_called_once()
|
|
assert runner._record_action_step.call_args.args[2] is result
|
|
assert turn.step_keys == ["recorded-file-step"]
|
|
assert turn.actions == [("android_file", status, result.text_summary)]
|
|
assert messages[-1].content == result.text_summary
|