89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from artemis.exploration.budget import ExplorationBudget
|
|
from artemis.runtime import task_control
|
|
|
|
|
|
def test_budget_excludes_demonstration_and_waiting_and_resets_only_on_continue(monkeypatch):
|
|
now = [0.0]
|
|
monkeypatch.setattr("artemis.runtime.task_limits.time.monotonic", lambda: now[0])
|
|
pause = Mock()
|
|
monkeypatch.setattr("artemis.exploration.budget.request_pause", pause)
|
|
repository = Mock()
|
|
budget = ExplorationBudget("run", 30, repository)
|
|
now[0] = 10
|
|
assert budget.update({"status": "demonstrating"}) == 20
|
|
now[0] = 500
|
|
assert budget.update({"status": "task_paused"}) == 20
|
|
now[0] = 1000
|
|
assert budget.update({"status": "running"}) == 20
|
|
now[0] = 1020
|
|
assert budget.update({"status": "running"}) == 0
|
|
pause.assert_called_once_with("run", reason="execution_budget")
|
|
now[0] = 1030
|
|
budget.update({"status": "task_paused"})
|
|
budget.update({"status": "task_paused"})
|
|
assert pause.call_count == 1
|
|
now[0] = 2000
|
|
assert budget.update({"status": "running", "budget_epoch": 1}) == 30
|
|
assert budget.pause_requested is False
|
|
assert budget.previous_elapsed == 40
|
|
now[0] = 2010
|
|
assert budget.update({"status": "running", "budget_epoch": 1}) == 20
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_answer_received_before_pause_acknowledgement_is_not_lost(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
|
|
class Lock:
|
|
device_id = "test-device"
|
|
description = "test"
|
|
max_concurrency = None
|
|
concurrency_mode = "per_device"
|
|
ingress = "test"
|
|
lock_scope = None
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.released = False
|
|
|
|
def acquire(self, **kwargs):
|
|
return None
|
|
|
|
def release(self):
|
|
self.released = True
|
|
|
|
monkeypatch.setattr(task_control, "DeviceExecutionLock", Lock)
|
|
control = task_control.TaskControl("run", Lock())
|
|
task_control.request_pause("run", reason="question")
|
|
task_control.add_instruction("run", "用户已回答", instruction_id="answer")
|
|
assert await asyncio.wait_for(control.checkpoint(), 1)
|
|
assert task_control.read_state("run")["status"] == "running"
|
|
assert control.budget_epoch == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_budget_resume_increments_epoch_after_lease_is_reacquired(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
old = SimpleNamespace(
|
|
device_id="test",
|
|
description="test",
|
|
max_concurrency=None,
|
|
concurrency_mode="per_device",
|
|
ingress="test",
|
|
lock_scope=None,
|
|
release=Mock(),
|
|
)
|
|
lock = SimpleNamespace(acquire=Mock(), release=Mock())
|
|
monkeypatch.setattr(task_control, "DeviceExecutionLock", Mock(return_value=lock))
|
|
control = task_control.TaskControl("run", old)
|
|
task_control.request_pause("run", reason="execution_budget")
|
|
task_control.request_resume("run")
|
|
assert await control.checkpoint()
|
|
lock.acquire.assert_called_once()
|
|
assert task_control.read_state("run")["budget_epoch"] == 1
|