chiwanko11
5d2283deaa
Introduce complete test coverage for the ARTEMIS platform across multiple layers: Unit tests: Add tests for agent nodes (Planner, Operator, Validator, Diagnoser, FlashRunner, Explorer), tools, utils (OCR/XML fusion, UI hierarchy), and data engine. Integration tests: Add live tool verification, reactive runner, IME, and standalone workflow tests. E2E tests: Add end-to-end scenario tests for media playback and failure recovery flows. Fixtures & mocks: Add mock inputs (audio, video, screenshot) and test utilities.
113 lines
4.0 KiB
Python
113 lines
4.0 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 asyncio
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from artemis.agents.operator.operator import OperatorNode
|
|
from artemis.context import ArtemisContext
|
|
from artemis.graph.graph import execution_check_node
|
|
from artemis.graph.state import State
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_snapshot_rollback_flow(tmp_path):
|
|
# Setup dummy directories
|
|
base_dir = tmp_path / "trace"
|
|
notes_dir = base_dir / "notes"
|
|
notes_dir.mkdir(parents=True)
|
|
|
|
task_plan_path = notes_dir / "task_plan.md"
|
|
task_plan_path.write_text("Initial Plan", encoding="utf-8")
|
|
|
|
# Mock Context
|
|
mock_ctx = MagicMock(spec=ArtemisContext)
|
|
mock_ctx.data_engine = MagicMock()
|
|
mock_ctx.data_engine.base_dir = base_dir
|
|
mock_ctx.execution_setup = None
|
|
|
|
# Mock State
|
|
mock_state = MagicMock(spec=State)
|
|
mock_state.subagent_calls = []
|
|
mock_state.asanitize_update = AsyncMock(return_value={"status": "success"})
|
|
mock_state.initial_goal = "Test goal"
|
|
mock_state.short_term_memory = ""
|
|
mock_state.operator_raw_data = {
|
|
"screenshot_b64": "YmFzZTY0",
|
|
"xml_hierarchy": "<hierarchy></hierarchy>",
|
|
"ocr_results": [],
|
|
}
|
|
mock_state.current_step_id = "test-step-id"
|
|
|
|
# Mock Controller and LLM to let Operator run minimally
|
|
mock_controller = MagicMock()
|
|
mock_controller.take_screenshot = AsyncMock(return_value="YmFzZTY0")
|
|
mock_device_data = MagicMock()
|
|
mock_device_data.elements = []
|
|
mock_controller.get_screen_data = AsyncMock(return_value=mock_device_data)
|
|
|
|
mock_llm = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.tool_calls = []
|
|
|
|
async def mock_astream(*args, **kwargs):
|
|
yield mock_response
|
|
|
|
mock_llm.astream.side_effect = mock_astream
|
|
mock_llm.bind_tools.return_value = mock_llm
|
|
|
|
# 1. Run Operator to trigger snapshot
|
|
with patch("artemis.agents.operator.operator.get_llm", return_value=mock_llm):
|
|
node = OperatorNode(mock_ctx)
|
|
await node(mock_state)
|
|
|
|
# Verify snapshot was created
|
|
snapshot_dir = notes_dir.with_name("notes_snapshot")
|
|
assert snapshot_dir.exists()
|
|
assert (snapshot_dir / "task_plan.md").read_text(encoding="utf-8") == "Initial Plan"
|
|
assert mock_ctx.task_plan_snapshot == snapshot_dir
|
|
|
|
# 2. Simulate Operator optimistic modification
|
|
task_plan_path.write_text("Modified Plan", encoding="utf-8")
|
|
|
|
# 3. Run execution_check_node with failure to trigger rollback
|
|
future = asyncio.Future()
|
|
future.set_result({"status": "failed"})
|
|
mock_ctx.checker_task = future
|
|
|
|
await execution_check_node(mock_state, mock_ctx)
|
|
|
|
# Verify rollback restored the file
|
|
assert task_plan_path.read_text(encoding="utf-8") == "Initial Plan"
|
|
|
|
# 4. Run Operator again to create new snapshot (simulate next attempt)
|
|
# First modify it again to simulate operator working on restored plan
|
|
task_plan_path.write_text("Modified Plan 2", encoding="utf-8")
|
|
|
|
with patch("artemis.agents.operator.operator.get_llm", return_value=mock_llm):
|
|
await node(mock_state)
|
|
|
|
assert (snapshot_dir / "task_plan.md").read_text(encoding="utf-8") == "Modified Plan 2"
|
|
|
|
# 5. Run execution_check_node with success to trigger commit (cleanup)
|
|
future_success = asyncio.Future()
|
|
future_success.set_result({"status": "success"})
|
|
mock_ctx.checker_task = future_success
|
|
|
|
await execution_check_node(mock_state, mock_ctx)
|
|
|
|
# Verify snapshot was deleted
|
|
assert not snapshot_dir.exists()
|