"""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"