tests/test_capabilities.py

import pytest

from app.capabilities import (
    normalize_capability, validate_capability, register_capabilities,
    get_capabilities, find_workers_with_capability, capability_to_endpoint_path,
    CAPABILITY_NAMES,
)
from app.models import JobType, WorkerRuntimeType


def test_valid_capability_names():
    names = list(CAPABILITY_NAMES.values())
    assert "iphone.text.embedding.private" in names
    assert "iphone.image.classify.local" in names
    assert "iphone.privacy.redact.local" in names


def test_invalid_capability_rejected():
    from app.capabilities import validate_capability
    cap = normalize_capability({"capability_name": "invalid.name", "runtime_type": WorkerRuntimeType.SAFARI_WASM})
    assert validate_capability(cap) is False


def test_capability_maps_to_endpoint():
    path = capability_to_endpoint_path("sess_abc", "iphone.text.embedding.private")
    assert path == "/api/session/sess_abc/cap/iphone.text.embedding.private"


def test_find_worker_for_capability():
    session_id = "sess_test"
    worker_id = "wk_test"
    caps = [{"capability_name": "iphone.text.embedding.private", "runtime_type": "safari_wasm", "model_id": "emb"}]
    register_capabilities(session_id, worker_id, caps)
    workers = find_workers_with_capability(session_id, JobType.TEXT_EMBEDDING)
    assert worker_id in workers