Files
agentic-mobile-control/tests/test_task_loop.py
T

79 lines
2.4 KiB
Python

from __future__ import annotations
from core.models import Bounds, Scene, SceneElement, Task
from runtime.executor import Executor, ExecutorConfig
from runtime.planner import PlannedStep, Planner
from runtime.task import TaskRunner, TaskRunnerConfig
from storage.artifact_store import ArtifactStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline
from tests.fakes import PNG_10X20
class ScriptedPlanner(Planner):
def __init__(self, steps: list[PlannedStep]) -> None:
self.steps = steps
def plan(self, *, goal, scene, context):
if len(context.step_results) >= len(self.steps):
return []
return [self.steps[len(context.step_results)]]
def goal_reached(self, *, goal, scene, context):
return len(context.step_results) >= len(self.steps) and all(
result.success for result in context.step_results
)
def test_task_runner_executes_loop_and_writes_timeline(tmp_path) -> None:
scene = Scene(
width=10,
height=20,
elements=[
SceneElement(
id="search",
type="input",
text="Search",
bounds=Bounds(1, 2, 4, 4),
)
],
)
planner = ScriptedPlanner(
[
PlannedStep(action="tap", description="tap search", args={"x": 3, "y": 4}),
PlannedStep(
action="input_text",
description="type query",
args={"text": "Mac mini"},
),
]
)
executor = Executor(
tools={
"tap": lambda **kwargs: {"ok": True, **kwargs},
"input_text": lambda **kwargs: {"ok": True, **kwargs},
},
config=ExecutorConfig(max_retries=1, backoff_seconds=0),
)
metadata = TaskMetadataStore(tmp_path / "tasks.sqlite3")
timeline = Timeline(ArtifactStore(tmp_path / "history"))
task = Task(goal="open app and search", device_id="iphone-1")
metadata.create_task(task)
runner = TaskRunner(
planner=planner,
executor=executor,
metadata_store=metadata,
timeline=timeline,
config=TaskRunnerConfig(max_steps=5),
observer=lambda device_id: scene,
screenshot_provider=lambda device_id: PNG_10X20,
)
result = runner.run(task)
assert result.status == "completed"
assert len(timeline.read(task.id)) == 2
assert metadata.get_task(task.id)["status"] == "completed"