182 lines
6.7 KiB
Python
182 lines
6.7 KiB
Python
"""Paused tasks yield their device and resume with a fresh lease."""
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
import artemis.runtime.task_control as task_control
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_releases_device_and_resume_waits_for_new_acquisition(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
events = []
|
|
|
|
class Lock:
|
|
def __init__(self, device_id, **kwargs):
|
|
self.device_id = device_id
|
|
self.description = kwargs.get("description", "test")
|
|
self.max_concurrency = kwargs.get("max_concurrency")
|
|
self.concurrency_mode = kwargs.get("concurrency_mode", "per_device")
|
|
self.ingress = kwargs.get("ingress", "web")
|
|
self.lock_scope = kwargs.get("lock_scope", "test")
|
|
|
|
def release(self):
|
|
events.append("released")
|
|
|
|
def acquire(self, cancel_event=None):
|
|
events.append("acquired")
|
|
|
|
monkeypatch.setattr(task_control, "DeviceExecutionLock", Lock)
|
|
original = Lock("test-device")
|
|
control = task_control.TaskControl("task-123", original)
|
|
task_control.request_pause("task-123")
|
|
waiter = asyncio.create_task(
|
|
control.checkpoint(
|
|
on_pause=lambda: _event("capture_stopped", events),
|
|
on_resume=lambda: _event("capture_started", events),
|
|
)
|
|
)
|
|
for _ in range(30):
|
|
if task_control.read_state("task-123").get("status") == "task_paused":
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
assert task_control.read_state("task-123")["status"] == "task_paused"
|
|
assert events == ["capture_stopped", "released"]
|
|
task_control.request_resume("task-123")
|
|
assert await asyncio.wait_for(waiter, 2) is True
|
|
assert events == ["capture_stopped", "released", "acquired", "capture_started"]
|
|
control.release()
|
|
assert task_control.read_state("task-123") == {}
|
|
|
|
|
|
async def _event(name, events):
|
|
events.append(name)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_interrupts_long_helper_and_discards_stale_result(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
events = []
|
|
started = asyncio.Event()
|
|
|
|
class Lock:
|
|
def __init__(self, device_id, **kwargs):
|
|
self.device_id = device_id
|
|
self.description = "test"
|
|
self.max_concurrency = None
|
|
self.concurrency_mode = "per_device"
|
|
self.ingress = "web"
|
|
self.lock_scope = "test"
|
|
|
|
def release(self):
|
|
events.append("released")
|
|
|
|
def acquire(self, cancel_event=None):
|
|
events.append("acquired")
|
|
|
|
monkeypatch.setattr(task_control, "DeviceExecutionLock", Lock)
|
|
control = task_control.TaskControl("task-123", Lock("device"))
|
|
|
|
async def diagnostic():
|
|
started.set()
|
|
try:
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
events.append("diagnostic_stopped")
|
|
|
|
result = asyncio.create_task(
|
|
control.run_interruptibly(
|
|
diagnostic(),
|
|
on_pause=lambda: _event("capture_stopped", events),
|
|
on_resume=lambda: _event("capture_started", events),
|
|
)
|
|
)
|
|
await asyncio.wait_for(started.wait(), 1)
|
|
task_control.request_pause("task-123")
|
|
for _ in range(30):
|
|
if task_control.read_state("task-123").get("status") == "task_paused":
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
assert task_control.read_state("task-123")["status"] == "task_paused"
|
|
assert events == ["diagnostic_stopped", "capture_stopped", "released"]
|
|
task_control.request_resume("task-123")
|
|
assert await asyncio.wait_for(result, 2) == (True, None)
|
|
assert events[-2:] == ["acquired", "capture_started"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_interruptible_call_returns_normal_result_without_pause(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
control = task_control.TaskControl("task-123", object())
|
|
|
|
async def model_call():
|
|
return "current result"
|
|
|
|
assert await control.run_interruptibly(model_call()) == (False, "current result")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_new_guidance_cancels_stale_model_without_releasing_device(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
released = []
|
|
|
|
class Lock:
|
|
def release(self):
|
|
released.append(True)
|
|
|
|
control = task_control.TaskControl("task-123", Lock())
|
|
started = asyncio.Event()
|
|
cancelled = asyncio.Event()
|
|
|
|
async def stale_model():
|
|
started.set()
|
|
try:
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
cancelled.set()
|
|
|
|
call = asyncio.create_task(control.run_interruptibly(stale_model()))
|
|
await asyncio.wait_for(started.wait(), 1)
|
|
task_control.add_instruction("task-123", "改为检查当前页面", source="mcp")
|
|
assert await asyncio.wait_for(call, 1) == (True, None)
|
|
assert cancelled.is_set()
|
|
assert released == []
|
|
assert task_control.consume_instructions(control.directory)["instruction"] == "改为检查当前页面"
|
|
|
|
|
|
def test_guidance_history_distinguishes_submission_from_runner_consumption(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
first = task_control.add_instruction("task-123", "从快捷设置进入", source="web")
|
|
second = task_control.add_instruction("task-123", "不要更改设置", source="mcp")
|
|
|
|
assert [item["status"] for item in task_control.list_instructions("task-123")] == [
|
|
"pending",
|
|
"pending",
|
|
]
|
|
consumed = task_control.consume_instructions(task_control.control_dir("task-123"))
|
|
assert consumed["instruction"] == "从快捷设置进入\n\n不要更改设置"
|
|
assert task_control.consume_instructions(task_control.control_dir("task-123")) is None
|
|
history = task_control.list_instructions("task-123")
|
|
assert [item["id"] for item in history] == [first["id"], second["id"]]
|
|
assert all(item["status"] == "consumed" and item["consumed_at"] for item in history)
|
|
|
|
|
|
def test_guidance_history_endpoint_returns_pending_and_consumed(tmp_path, monkeypatch):
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from apps.admin_console.routers.tasks import router
|
|
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
task_control.add_instruction("task-123", "从快捷设置进入")
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
client = TestClient(app)
|
|
response = client.get("/api/tasks/task-123/guidance")
|
|
assert response.status_code == 200
|
|
assert response.json()[0]["status"] == "pending"
|
|
task_control.consume_instructions(task_control.control_dir("task-123"))
|
|
response = client.get("/api/tasks/task-123/guidance")
|
|
assert response.json()[0]["status"] == "consumed"
|