Aether/tests/services/test_pool_preset_dimensions.py

"""Tests for pool preset dimension registry and built-in dimensions."""

from __future__ import annotations

from types import SimpleNamespace

import src.services.provider.pool.dimensions  # noqa: F401
from src.services.provider.pool.dimensions import get_preset_dimension, get_preset_names


def _key(metadata: dict, *, plan_type: str | None = None) -> SimpleNamespace:
    return SimpleNamespace(upstream_metadata=metadata, oauth_plan_type=plan_type)


def test_registry_discovers_builtin_dimensions() -> None:
    names = get_preset_names()
    assert {
        "free_team_first",
        "recent_refresh",
        "quota_balanced",
        "single_account",
        "priority_first",
        "health_first",
        "latency_first",
        "cost_first",
    }.issubset(names)


def test_universal_dimensions_are_applicable_to_any_provider() -> None:
    for name in ("quota_balanced", "single_account"):
        dim = get_preset_dimension(name)
        assert dim is not None
        assert dim.is_applicable("openai") is True
        assert dim.is_applicable("codex") is True
        assert dim.is_applicable("unknown_provider") is True


def test_provider_specific_dimensions_are_filtered() -> None:
    for name in ("free_team_first", "recent_refresh"):
        dim = get_preset_dimension(name)
        assert dim is not None
        assert dim.is_applicable("codex") is True
        assert dim.is_applicable("kiro") is True
        assert dim.is_applicable("openai") is False


def test_builtin_dimensions_compute_metric_in_range() -> None:
    all_key_ids = ["k1", "k2", "k3"]
    lru_scores = {"k1": 100.0, "k2": 800.0, "k3": 300.0}
    keys_by_id = {
        "k1": _key(
            {
                "codex": {
                    "plan_type": "plus",
                    "primary_reset_seconds": 900,
                    "primary_used_percent": 60,
                }
            }
        ),
        "k2": _key(
            {
                "codex": {
                    "plan_type": "free",
                    "primary_reset_seconds": 120,
                    "primary_used_percent": 20,
                }
            }
        ),
        "k3": _key(
            {
                "kiro": {
                    "next_reset_at": 4102444800,
                    "usage_percentage": 45,
                    "subscription_title": "Kiro Team",
                }
            },
            plan_type="team",
        ),
    }

    for name in get_preset_names():
        dim = get_preset_dimension(name)
        assert dim is not None
        mode = dim.default_mode
        metric = dim.compute_metric(
            key_id="k1",
            all_key_ids=all_key_ids,
            keys_by_id=keys_by_id,
            lru_scores=lru_scores,
            context={},
            mode=mode,
        )
        assert 0.0 <= metric <= 1.0