Aether/tests/unit/test_provider_candidate_ordering.py

from __future__ import annotations

from types import SimpleNamespace

from src.services.scheduling.aware_scheduler import ProviderCandidate


def _make_candidate(
    *,
    provider_priority: int | None,
    internal_priority: int | None,
    provider_id: str,
    endpoint_id: str,
    key_id: str,
) -> ProviderCandidate:
    provider = SimpleNamespace(id=provider_id, name="p", provider_priority=provider_priority)
    endpoint = SimpleNamespace(id=endpoint_id)
    key = SimpleNamespace(id=key_id, internal_priority=internal_priority)
    return ProviderCandidate(provider=provider, endpoint=endpoint, key=key)  # type: ignore[arg-type]


def test_provider_candidate_is_orderable() -> None:
    c1 = _make_candidate(
        provider_priority=1,
        internal_priority=1,
        provider_id="p1",
        endpoint_id="e1",
        key_id="k1",
    )
    c2 = _make_candidate(
        provider_priority=1,
        internal_priority=2,
        provider_id="p1",
        endpoint_id="e1",
        key_id="k2",
    )

    assert c1 < c2
    assert sorted([c2, c1]) == [c1, c2]


def test_provider_candidate_can_break_ties_in_tuple_sort() -> None:
    c1 = _make_candidate(
        provider_priority=1,
        internal_priority=1,
        provider_id="p1",
        endpoint_id="e1",
        key_id="k1",
    )
    c2 = _make_candidate(
        provider_priority=1,
        internal_priority=2,
        provider_id="p1",
        endpoint_id="e1",
        key_id="k2",
    )

    # When the first tuple element ties, Python will compare the second element.
    # This should not raise:
    # TypeError: '<' not supported between instances of 'ProviderCandidate' and 'ProviderCandidate'
    pairs = [(0, c2), (0, c1)]
    assert [c for _score, c in sorted(pairs)] == [c1, c2]


def test_provider_candidate_none_priorities_sort_last() -> None:
    c1 = _make_candidate(
        provider_priority=1,
        internal_priority=None,
        provider_id="p1",
        endpoint_id="e1",
        key_id="k1",
    )
    c2 = _make_candidate(
        provider_priority=None,
        internal_priority=None,
        provider_id="p2",
        endpoint_id="e1",
        key_id="k2",
    )

    assert sorted([c2, c1]) == [c1, c2]