cloud

tests/test_task_scheduler.py

@@ -0,0 +1,219 @@
+"""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 (
+    AssignmentStrategy,
+    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"
+
+
+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"