tests/test_capability_resolver.py

import asyncio
import sys
import time
import types
import unittest
from enum import IntEnum
from pathlib import Path

from pydantic import ValidationError

BACKEND_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(BACKEND_ROOT))


def _install_import_stubs() -> None:
    auth_guard = types.ModuleType("api.auth_guard")

    class AuthRole(IntEnum):
        MACHINE = 1

    auth_guard.AuthRole = AuthRole
    auth_guard.require_role = lambda _role: (lambda: None)
    sys.modules["api.auth_guard"] = auth_guard

    health_module = types.ModuleType("api.health_manager")

    class HealthyManager:
        async def is_healthy(self, _worker_id: str) -> bool:
            return True

    health_module.health_manager = HealthyManager()
    sys.modules["api.health_manager"] = health_module


_install_import_stubs()

from api.marketplace import WORKERS_REGISTRY, WorkerCapability
from api.resolver import CapabilityResolver, ResolverConstraints


class CapabilityResolverTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self._original_registry = dict(WORKERS_REGISTRY)
        WORKERS_REGISTRY.clear()

    def tearDown(self):
        WORKERS_REGISTRY.clear()
        WORKERS_REGISTRY.update(self._original_registry)

    def register_worker(self, **overrides) -> WorkerCapability:
        values = {
            "id": "worker",
            "name": "Worker",
            "version": "1.0.0",
            "last_seen": int(time.time()),
            "capabilities": ["vision"],
            "cost": 1.0,
            "latency": 100.0,
            "region": "global",
            "gpu": False,
            "priority": 10,
        }
        values.update(overrides)
        worker = WorkerCapability(**values)
        WORKERS_REGISTRY[worker.id] = worker
        return worker

    async def test_uses_semver_not_lexicographic_ordering(self):
        self.register_worker(id="v2", version="2.0.0", priority=20)
        self.register_worker(id="v10", version="10.0.0", priority=10)

        worker = await CapabilityResolver.resolve(
            "vision",
            ResolverConstraints(min_version="2.0.0"),
        )

        self.assertIsNotNone(worker)
        self.assertEqual(worker.id, "v10")

    async def test_excludes_prerelease_below_stable_minimum(self):
        self.register_worker(id="candidate", version="2.0.0-rc.1", priority=1)
        self.register_worker(id="stable", version="2.0.0", priority=10)

        worker = await CapabilityResolver.resolve(
            "vision",
            ResolverConstraints(min_version="2.0.0"),
        )

        self.assertIsNotNone(worker)
        self.assertEqual(worker.id, "stable")

    def test_rejects_malformed_minimum_semver(self):
        with self.assertRaises(ValidationError):
            ResolverConstraints(min_version="2.0")

    async def test_excludes_malformed_worker_version_only_when_constrained(self):
        self.register_worker(id="malformed", version="not-a-version", priority=1)
        self.register_worker(id="valid", version="2.0.0", priority=10)

        worker = await CapabilityResolver.resolve(
            "vision",
            ResolverConstraints(min_version="2.0.0"),
        )

        self.assertIsNotNone(worker)
        self.assertEqual(worker.id, "valid")

    async def test_prefers_requested_region_when_available(self):
        self.register_worker(id="global", region="global", priority=1)
        self.register_worker(id="eu", region="eu-west", priority=20)

        worker = await CapabilityResolver.resolve(
            "vision",
            ResolverConstraints(preferred_region="eu-west"),
        )

        self.assertIsNotNone(worker)
        self.assertEqual(worker.id, "eu")

    async def test_falls_back_when_requested_region_is_unavailable(self):
        self.register_worker(id="global", region="global", priority=1)

        worker = await CapabilityResolver.resolve(
            "vision",
            ResolverConstraints(preferred_region="eu-west"),
        )

        self.assertIsNotNone(worker)
        self.assertEqual(worker.id, "global")

    async def test_applies_priority_threshold_with_lower_values_preferred(self):
        self.register_worker(id="allowed", priority=100, cost=10.0)
        self.register_worker(id="excluded", priority=101, cost=0.0)

        worker = await CapabilityResolver.resolve(
            "vision",
            ResolverConstraints(min_priority=100),
        )

        self.assertIsNotNone(worker)
        self.assertEqual(worker.id, "allowed")

    async def test_preserves_existing_cost_latency_and_gpu_constraints(self):
        self.register_worker(
            id="cpu",
            cost=1.0,
            latency=50.0,
            gpu=False,
            priority=1,
        )
        self.register_worker(
            id="gpu",
            cost=2.0,
            latency=100.0,
            gpu=True,
            priority=10,
        )

        worker = await CapabilityResolver.resolve(
            "vision",
            ResolverConstraints(
                max_cost=2.0,
                max_latency=100.0,
                require_gpu=True,
            ),
        )

        self.assertIsNotNone(worker)
        self.assertEqual(worker.id, "gpu")


if __name__ == "__main__":
    unittest.main()