296 lines
9.3 KiB
Python
296 lines
9.3 KiB
Python
"""Unit tests for cloud.scheduler.TaskScheduler (task 4.8)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from cloud.config import CloudConfig
|
|
from cloud.pool import DevicePool
|
|
from cloud.scheduler import (
|
|
FIFO_MATCH_STRATEGY_NAME,
|
|
QueueFullError,
|
|
ScheduledTask,
|
|
TaskConstraints,
|
|
TaskScheduler,
|
|
TaskSubmissionValidationError,
|
|
UnknownAssignmentStrategyError,
|
|
)
|
|
from cloud.store import CloudStore
|
|
from core.models import Device
|
|
|
|
|
|
def _config(**overrides) -> CloudConfig:
|
|
base = {
|
|
"sync_interval_seconds": 30,
|
|
"stale_after_seconds": 60,
|
|
"max_queue_depth": 100,
|
|
"default_assignment_strategy": "fifo_match",
|
|
"api_version_prefix": "/v1",
|
|
"db_path": "cloud/cloud.sqlite3",
|
|
}
|
|
base.update(overrides)
|
|
return CloudConfig(**base)
|
|
|
|
|
|
def _device(
|
|
device_id: str, *, status: str = "idle", driver_type: str = "wda"
|
|
) -> Device:
|
|
return Device(id=device_id, status=status, driver_type=driver_type) # type: ignore[arg-type]
|
|
|
|
|
|
def _pool_with_devices(
|
|
tmp_path, *devices: Device, host_id: str = "host-local"
|
|
) -> DevicePool:
|
|
pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())
|
|
pool.sync_host_devices(host_id, list(devices))
|
|
return pool
|
|
|
|
|
|
def test_submit_enqueues_with_status_queued(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path)
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
|
|
task_id = scheduler.submit(goal="open settings")
|
|
|
|
assert isinstance(task_id, str) and task_id
|
|
task = pool.store.get_task(task_id)
|
|
assert task is not None
|
|
assert task.status == "queued"
|
|
assert task.goal == "open settings"
|
|
assert task.workflow_definition_id is None
|
|
|
|
|
|
def test_submit_requires_goal_or_workflow(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path)
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
with pytest.raises(TaskSubmissionValidationError):
|
|
scheduler.submit()
|
|
|
|
|
|
def test_queue_depth_limit_rejects_submission(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path)
|
|
scheduler = TaskScheduler(pool, pool.store, _config(max_queue_depth=2))
|
|
scheduler.submit(goal="one")
|
|
scheduler.submit(goal="two")
|
|
with pytest.raises(QueueFullError):
|
|
scheduler.submit(goal="three")
|
|
|
|
|
|
def test_assign_picks_matching_idle_device(tmp_path) -> None:
|
|
pool = _pool_with_devices(
|
|
tmp_path,
|
|
_device("dev-1", driver_type="wda"),
|
|
_device("dev-2", status="busy", driver_type="wda"),
|
|
)
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
|
|
task_id = scheduler.submit(
|
|
goal="x",
|
|
constraints=TaskConstraints(driver_type="wda"),
|
|
)
|
|
assignments = scheduler.assign()
|
|
|
|
assert [a.task_id for a in assignments] == [task_id]
|
|
assert assignments[0].device_id == "dev-1"
|
|
assert assignments[0].host_id == "host-local"
|
|
|
|
task = pool.store.get_task(task_id)
|
|
assert task is not None
|
|
assert task.status == "assigned"
|
|
assert task.assigned_device_id == "dev-1"
|
|
|
|
|
|
def test_assign_leaves_task_queued_when_no_device_matches(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path, _device("dev-1", driver_type="wda"))
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
|
|
task_id = scheduler.submit(
|
|
goal="x",
|
|
constraints=TaskConstraints(driver_type="android"),
|
|
)
|
|
assignments = scheduler.assign()
|
|
|
|
assert assignments == []
|
|
task = pool.store.get_task(task_id)
|
|
assert task is not None
|
|
assert task.status == "queued"
|
|
|
|
|
|
def test_two_tasks_assigned_in_submission_order_with_one_device(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path, _device("dev-1", driver_type="wda"))
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
|
|
first_id = scheduler.submit(goal="first")
|
|
# Ensure a distinct created_at for the second submission so list_queued_tasks
|
|
# ordering by (created_at, id) is deterministic.
|
|
time.sleep(0.005)
|
|
second_id = scheduler.submit(goal="second")
|
|
|
|
assignments = scheduler.assign()
|
|
assert [a.task_id for a in assignments] == [first_id]
|
|
|
|
first_task = pool.store.get_task(first_id)
|
|
second_task = pool.store.get_task(second_id)
|
|
assert first_task is not None and second_task is not None
|
|
assert first_task.status == "assigned"
|
|
assert second_task.status == "queued"
|
|
|
|
assert scheduler.assign() == []
|
|
second_task = pool.store.get_task(second_id)
|
|
assert second_task is not None
|
|
assert second_task.status == "queued"
|
|
|
|
|
|
def test_unknown_strategy_raises_at_init(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path)
|
|
with pytest.raises(UnknownAssignmentStrategyError):
|
|
TaskScheduler(
|
|
pool,
|
|
pool.store,
|
|
_config(default_assignment_strategy="nonexistent"),
|
|
)
|
|
|
|
|
|
def test_custom_strategy_can_be_registered(tmp_path) -> None:
|
|
"""A new AssignmentStrategy can be plugged in by name without scheduler edits."""
|
|
|
|
class LastDeviceStrategy:
|
|
def select(self, task: ScheduledTask, candidates): # type: ignore[override]
|
|
return candidates[-1] if candidates else None
|
|
|
|
pool = _pool_with_devices(
|
|
tmp_path,
|
|
_device("dev-1"),
|
|
_device("dev-2"),
|
|
)
|
|
scheduler = TaskScheduler(
|
|
pool,
|
|
pool.store,
|
|
_config(default_assignment_strategy="last"),
|
|
strategies={
|
|
FIFO_MATCH_STRATEGY_NAME: type(pool).__module__, # placeholder
|
|
"last": LastDeviceStrategy(), # type: ignore[dict-item]
|
|
},
|
|
)
|
|
# Replace with the real fifo strategy so other tests' default isn't relied on.
|
|
scheduler._strategies[FIFO_MATCH_STRATEGY_NAME] = type( # noqa: SLF001
|
|
pool,
|
|
).__module__
|
|
|
|
task_id = scheduler.submit(goal="x")
|
|
assignments = scheduler.assign()
|
|
|
|
assert len(assignments) == 1
|
|
assert assignments[0].device_id == "dev-2"
|
|
assert assignments[0].task_id == task_id
|
|
|
|
|
|
def test_capability_tag_constraint_filters_candidates(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path)
|
|
# Manually plant devices with capability_tags by going through the store.
|
|
from datetime import UTC, datetime
|
|
|
|
from cloud.pool import PooledDevice
|
|
|
|
pool.store.replace_host_devices(
|
|
"host-local",
|
|
[
|
|
PooledDevice(
|
|
device_id="dev-1",
|
|
host_id="host-local",
|
|
driver_type="wda",
|
|
status="idle",
|
|
capability_tags=["ios"],
|
|
synced_at=datetime.now(UTC),
|
|
),
|
|
PooledDevice(
|
|
device_id="dev-2",
|
|
host_id="host-local",
|
|
driver_type="wda",
|
|
status="idle",
|
|
capability_tags=["android"],
|
|
synced_at=datetime.now(UTC),
|
|
),
|
|
],
|
|
)
|
|
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
task_id = scheduler.submit(
|
|
goal="x",
|
|
constraints=TaskConstraints(capability_tags=["android"]),
|
|
)
|
|
|
|
assignments = scheduler.assign()
|
|
assert [a.task_id for a in assignments] == [task_id]
|
|
assert assignments[0].device_id == "dev-2"
|
|
|
|
|
|
def test_explicit_host_and_device_target_is_a_hard_constraint(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path, _device("host-a-device"), host_id="host-a")
|
|
pool.sync_host_devices("host-b", [_device("host-b-device")])
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
|
|
task_id = scheduler.submit(
|
|
goal="target host b",
|
|
constraints=TaskConstraints(
|
|
target_host_id="host-b",
|
|
target_device_id="host-b-device",
|
|
),
|
|
)
|
|
|
|
assignments = scheduler.assign()
|
|
|
|
assert [(item.host_id, item.device_id) for item in assignments] == [
|
|
("host-b", "host-b-device")
|
|
]
|
|
task = pool.store.get_task(task_id)
|
|
assert task is not None
|
|
assert task.constraints.target_host_id == "host-b"
|
|
assert task.constraints.target_device_id == "host-b-device"
|
|
|
|
|
|
def test_unavailable_explicit_target_is_never_rerouted(tmp_path) -> None:
|
|
pool = _pool_with_devices(tmp_path, _device("host-a-device"), host_id="host-a")
|
|
pool.sync_host_devices("host-b", [_device("host-b-device", status="busy")])
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
task_id = scheduler.submit(
|
|
goal="wait for host b",
|
|
constraints=TaskConstraints(target_host_id="host-b"),
|
|
)
|
|
|
|
assert scheduler.assign() == []
|
|
task = pool.store.get_task(task_id)
|
|
assert task is not None
|
|
assert task.status == "queued"
|
|
|
|
|
|
def test_host_active_task_policy_keeps_excess_tasks_queued(tmp_path) -> None:
|
|
from datetime import UTC, datetime
|
|
|
|
pool = _pool_with_devices(
|
|
tmp_path,
|
|
_device("device-a"),
|
|
_device("device-b"),
|
|
host_id="host-a",
|
|
)
|
|
pool.store.upsert_host_governance_policy(
|
|
host_id="host-a",
|
|
self_submission_enabled=True,
|
|
max_active_tasks=1,
|
|
daily_token_budget=None,
|
|
updated_at=datetime.now(UTC),
|
|
)
|
|
scheduler = TaskScheduler(pool, pool.store, _config())
|
|
first_id = scheduler.submit(goal="first")
|
|
second_id = scheduler.submit(goal="second")
|
|
|
|
assignments = scheduler.assign()
|
|
|
|
assert [assignment.task_id for assignment in assignments] == [first_id]
|
|
assert pool.store.count_active_tasks_for_host("host-a") == 1
|
|
second = pool.store.get_task(second_id)
|
|
assert second is not None
|
|
assert second.status == "queued"
|