tests/test_credits_election.py

"""KI-085 (2026-05-15) — proactive credit tracking + election gate.

Tests:
  1. update_credits_from_groq parser variants (valid / threshold / missing / malformed)
  2. NIM local rate-meter increments + 60s reset
  3. _is_election_eligible gates by credits_remaining vs credits_low_water
  4. elect_primary_and_backup integration: credit-exhausted candidates skipped

Run:
    .venv/bin/python -m unittest tests.test_credits_election -v
"""

from __future__ import annotations

import time
import unittest
from unittest import mock

from backend import llm_health
from backend.llm_health import (
    GROQ_TOKENS_LOW_WATER,
    ModelHealth,
    NIM_REQ_PER_MIN_LOW_WATER,
    _has_credits,
    _is_election_eligible,
    record_nim_call,
    update_credits_from_groq,
)


def _fresh_state():
    """Wipe in-memory state + mark loaded so tests don't trigger disk reads."""
    llm_health._STATE.clear()
    llm_health._STATE_LOADED = True
    llm_health._NIM_CALL_TIMES.clear()


def _healthy_now(model: str) -> ModelHealth:
    """Build a ModelHealth that passes every check EXCEPT credits — so
    flipping credits is the only failure surface in these tests."""
    h = ModelHealth(model=model)
    h.status = "healthy"
    h.latency_ms = 200
    h.tested_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    return h


class TestGroqHeaderParser(unittest.TestCase):
    def setUp(self) -> None:
        _fresh_state()

    def test_546_tokens_remaining_gates_out(self) -> None:
        """546 < 5000 low water → credit-gate fails."""
        model = "groq:llama-3.3-70b-versatile"
        update_credits_from_groq("brain", model, {
            "x-ratelimit-remaining-tokens-day": "546",
            "x-ratelimit-reset-tokens-day": "1h2m",
        })
        h = llm_health._STATE[model]
        self.assertEqual(h.credits_remaining, 546.0)
        self.assertEqual(h.credits_unit, "tokens_day")
        self.assertEqual(h.credits_low_water, GROQ_TOKENS_LOW_WATER)
        self.assertFalse(_has_credits(h, time.monotonic()))

    def test_50000_tokens_remaining_is_eligible(self) -> None:
        """50K >> 5K low water → credit-gate passes."""
        model = "groq:llama-3.3-70b-versatile"
        update_credits_from_groq("brain", model, {
            "x-ratelimit-remaining-tokens-day": "50000",
            "x-ratelimit-reset-tokens-day": "5h",
        })
        h = llm_health._STATE[model]
        self.assertEqual(h.credits_remaining, 50000.0)
        self.assertTrue(_has_credits(h, time.monotonic()))

    def test_missing_header_is_noop(self) -> None:
        """No daily-tokens header → state untouched (credits_remaining stays None)."""
        model = "groq:llama-3.3-70b-versatile"
        update_credits_from_groq("brain", model, {})
        self.assertNotIn(model, llm_health._STATE)

    def test_malformed_value_is_noop(self) -> None:
        """Garbage in the header → no-op + warning logged (not raised)."""
        model = "groq:llama-3.3-70b-versatile"
        update_credits_from_groq("brain", model, {
            "x-ratelimit-remaining-tokens-day": "not-a-number",
        })
        # State remains absent because update bailed before stamping.
        self.assertNotIn(model, llm_health._STATE)

    def test_reset_string_parses_duration(self) -> None:
        """`1h2m` → ~3720s offset; `45s` → ~45s offset; `30m` → ~1800s."""
        now_mono = time.monotonic()
        from backend.llm_health import _parse_reset_seconds
        v = _parse_reset_seconds("1h2m", now_mono)
        self.assertIsNotNone(v)
        self.assertAlmostEqual(v - now_mono, 3720, delta=2)
        v = _parse_reset_seconds("45s", now_mono)
        self.assertAlmostEqual(v - now_mono, 45, delta=2)
        v = _parse_reset_seconds("30m", now_mono)
        self.assertAlmostEqual(v - now_mono, 1800, delta=2)
        # Bare seconds-from-now
        v = _parse_reset_seconds("120", now_mono)
        self.assertAlmostEqual(v - now_mono, 120, delta=2)
        # Empty / None
        self.assertIsNone(_parse_reset_seconds("", now_mono))
        self.assertIsNone(_parse_reset_seconds(None, now_mono))


class TestNimRateMeter(unittest.TestCase):
    def setUp(self) -> None:
        _fresh_state()

    def test_30_calls_in_60s_eligible(self) -> None:
        """30 successful calls → 40-30=10 remaining, > 5 low water → eligible."""
        model = "qwen/qwen3-next-80b-a3b-instruct"
        for _ in range(30):
            record_nim_call("brain", model)
        h = llm_health._STATE[model]
        self.assertEqual(h.credits_remaining, 10.0)
        self.assertEqual(h.credits_unit, "requests_min")
        self.assertTrue(_has_credits(h, time.monotonic()))

    def test_36_calls_in_60s_gated_out(self) -> None:
        """36 successful calls → 40-36=4 remaining, < 5 low water → gated."""
        model = "qwen/qwen3-next-80b-a3b-instruct"
        for _ in range(36):
            record_nim_call("brain", model)
        h = llm_health._STATE[model]
        self.assertEqual(h.credits_remaining, 4.0)
        self.assertFalse(_has_credits(h, time.monotonic()))

    def test_60s_window_resets(self) -> None:
        """After 60s elapses, _has_credits flips back to True via stale-reset."""
        model = "qwen/qwen3-next-80b-a3b-instruct"
        for _ in range(36):
            record_nim_call("brain", model)
        h = llm_health._STATE[model]
        self.assertFalse(_has_credits(h, time.monotonic()))
        # Simulate 65 seconds passing: credits_reset_at is now in the past.
        future_now = (h.credits_reset_at or time.monotonic()) + 5.0
        self.assertTrue(_has_credits(h, future_now))


class TestElectionCreditGate(unittest.TestCase):
    """Integration: elect_primary skips quota-exhausted candidates."""

    def setUp(self) -> None:
        _fresh_state()

    def test_groq_below_water_skipped_in_election(self) -> None:
        """Groq has fast probe latency but is below daily-tokens water →
        elector falls through to next candidate."""
        groq_model = "groq:llama-3.3-70b-versatile"
        nim_model = "qwen/qwen3-next-80b-a3b-instruct"
        # Both healthy + fresh probe + low latency. Groq is faster (100ms).
        gh = _healthy_now(groq_model)
        gh.latency_ms = 100
        gh.credits_remaining = 100.0  # < 5000 low_water
        gh.credits_unit = "tokens_day"
        gh.credits_low_water = GROQ_TOKENS_LOW_WATER
        gh.credits_observed_at = time.monotonic()
        llm_health._STATE[groq_model] = gh

        nh = _healthy_now(nim_model)
        nh.latency_ms = 300
        llm_health._STATE[nim_model] = nh

        with mock.patch.object(
            llm_health, "_chain_for", return_value=[groq_model, nim_model]
        ):
            primary = llm_health.get_primary("brain")
        self.assertEqual(primary, nim_model,
                         "Quota-exhausted Groq should be skipped despite faster latency.")

    def test_nim_preferred_over_faster_groq_when_eligible(self) -> None:
        """KI-201 chain-order: chain head wins regardless of latency.

        Under KI-201 the chain-head model is always primary when eligible.
        Pre-KI-201 (KI-087) the elector applied a NIM-first preference on
        top of latency-score; that preference is now redundant because the
        canonical chains in nvidia_nim_llm.py list NIM models first by
        construction. This test exercises the new contract: with NIM as
        chain-head and Groq slower in the chain, NIM wins on POSITION
        not on latency-vs-NIM-preference."""
        groq_model = "groq:llama-3.3-70b-versatile"
        nim_model = "qwen/qwen3-next-80b-a3b-instruct"
        gh = _healthy_now(groq_model)
        gh.latency_ms = 100  # faster (irrelevant under KI-201)
        gh.credits_remaining = 10000.0  # well above water
        gh.credits_unit = "tokens_day"
        gh.credits_low_water = GROQ_TOKENS_LOW_WATER
        gh.credits_observed_at = time.monotonic()
        llm_health._STATE[groq_model] = gh

        nh = _healthy_now(nim_model)
        nh.latency_ms = 300  # slower but chain-head
        llm_health._STATE[nim_model] = nh

        # NIM is chain-head; under chain-order election it wins regardless
        # of Groq being faster.
        with mock.patch.object(
            llm_health, "_chain_for", return_value=[nim_model, groq_model]
        ):
            primary = llm_health.get_primary("brain")
        self.assertEqual(
            primary, nim_model,
            "KI-201: chain-head wins regardless of latency. NIM is listed "
            "first in BRAIN_CHAIN by construction (KI-175); chain-order "
            "election picks it without needing a separate NIM-first rule.",
        )

    def test_groq_picked_when_nim_pool_empty(self) -> None:
        """KI-201 chain-order fallthrough: when NO eligible NIM candidate
        exists (all NIM models down / out of credits / not in chain),
        election walks the chain in order and picks the FIRST eligible
        non-NIM candidate. Locks in the safety net so a full NIM regional
        outage still produces a working brain call.

        Pre-KI-201 (KI-087) this test asserted Groq won via latency-score
        even when listed second in the chain. Under chain-order election
        the chain is the truth, so we list Groq first to exercise the
        same fallthrough behaviour."""
        groq_model = "groq:llama-3.3-70b-versatile"
        or_model = "openrouter:openai/gpt-oss-120b"
        gh = _healthy_now(groq_model)
        gh.latency_ms = 100
        gh.credits_remaining = 10000.0
        gh.credits_unit = "tokens_day"
        gh.credits_low_water = GROQ_TOKENS_LOW_WATER
        gh.credits_observed_at = time.monotonic()
        llm_health._STATE[groq_model] = gh

        oh = _healthy_now(or_model)
        oh.latency_ms = 800  # slower (irrelevant under KI-201 chain-order)
        llm_health._STATE[or_model] = oh

        # Chain has NO NIM candidates; Groq is chain-head.
        with mock.patch.object(
            llm_health, "_chain_for", return_value=[groq_model, or_model]
        ):
            primary = llm_health.get_primary("brain")
        self.assertEqual(
            primary, groq_model,
            "KI-201 chain-order: no NIM eligible → election picks the "
            "FIRST eligible candidate in chain order (Groq is chain-head).",
        )

    def test_none_credits_is_permissive(self) -> None:
        """Cold-start: a candidate with credits_remaining=None must be electable."""
        model = "qwen/qwen3-next-80b-a3b-instruct"
        h = _healthy_now(model)
        # credits_remaining is None by default — leave it alone.
        llm_health._STATE[model] = h
        self.assertTrue(_is_election_eligible(h, time.monotonic()),
                        "Cold-start (None credits) must NOT gate out a healthy candidate.")


class TestStrictChainCreditExhausted(unittest.TestCase):
    """KI-122 (2026-05-15) — strict banner rule.

    Banner fires ONLY when zero chain members are available_for_calls AND at
    least one has a credit-exhaustion signal that's strictly in-window
    (credits_remaining<=low_water AND credits_reset_at is in the future).

    Pre-KI-122 holes the new rule must close:
      H1) one healthy NIM primary (credits=None) + one credit-low Groq backup
          → banner USED to fire because `any_signal=True` on Groq + skip on
            NIM's None signal. Must NOT fire.
      H2) every chain member has reset_at=None (e.g. OpenRouter wallet) and
          credits below water → banner USED to fire. Must NOT fire by itself
          (probe success is the authoritative truth; wallet-zero with
          successful probes means free-tier models).
      H3) every member is genuinely gated (low + future reset) AND not
          available → banner SHOULD fire.
    """

    def setUp(self) -> None:
        _fresh_state()
        # Import locally so test discovery doesn't load admin until needed.
        from backend.admin import _chain_summary
        self._chain_summary = _chain_summary

    def _build(self, chain_models, model_states):
        """Inject states + run _chain_summary against a mocked chain."""
        for m, h in model_states.items():
            llm_health._STATE[m] = h
        now_mono = time.monotonic()
        with mock.patch.object(
            llm_health, "_chain_for", return_value=list(chain_models),
        ):
            return self._chain_summary(
                "brain", {"brain": list(chain_models)}, llm_health._STATE, now_mono,
            )

    def test_healthy_nim_plus_exhausted_groq_does_not_fire_banner(self) -> None:
        """H1 — NIM primary HEALTHY/None-credits + Groq backup credit-low.
        Banner must STAY DOWN because NIM is available_for_calls."""
        nim = "qwen/qwen3-next-80b-a3b-instruct"
        groq = "groq:llama-3.3-70b-versatile"
        nh = _healthy_now(nim)
        nh.latency_ms = 300
        # NIM has NO credit signal (None) — cold start permissive.
        gh = _healthy_now(groq)
        gh.latency_ms = 100
        gh.credits_remaining = 100.0
        gh.credits_unit = "tokens_day"
        gh.credits_low_water = GROQ_TOKENS_LOW_WATER
        gh.credits_reset_at = time.monotonic() + 3600.0  # future reset
        gh.credits_observed_at = time.monotonic()
        block = self._build([nim, groq], {nim: nh, groq: gh})
        self.assertFalse(
            block["chain_credit_exhausted"],
            "Banner must NOT fire when a healthy cold-start NIM primary "
            "shares a chain with one credit-exhausted Groq backup.",
        )
        # Sanity: NIM should be in_use, Groq should be gated.
        names = {mi["model"]: mi for mi in block["chain_members"]}
        self.assertTrue(names[nim]["available_for_calls"])
        self.assertTrue(names[nim]["is_current_primary"])
        self.assertFalse(names[groq]["available_for_calls"])
        self.assertTrue(names[groq]["credit_exhausted"])

    def test_credits_reset_at_none_does_not_fire_banner(self) -> None:
        """H2 — every member has reset_at=None + credits below water (typical
        OpenRouter usd_balance free-tier shape). Banner must NOT fire on the
        reset_at=None signal alone; the strict rule requires an in-window reset."""
        m1 = "openrouter:openai/gpt-oss-120b"
        m2 = "openrouter:meta-llama/llama-3.3-70b-instruct:free"
        h1 = _healthy_now(m1)
        h1.latency_ms = 800
        h1.credits_remaining = 0.0
        h1.credits_unit = "usd_balance"
        h1.credits_low_water = 0.05
        h1.credits_reset_at = None  # OpenRouter wallet: no scheduled reset
        h1.credits_observed_at = time.monotonic()
        h2 = _healthy_now(m2)
        h2.latency_ms = 900
        h2.credits_remaining = 0.0
        h2.credits_unit = "usd_balance"
        h2.credits_low_water = 0.05
        h2.credits_reset_at = None
        h2.credits_observed_at = time.monotonic()
        block = self._build([m1, m2], {m1: h1, m2: h2})
        self.assertFalse(
            block["chain_credit_exhausted"],
            "Banner must NOT fire when credits_reset_at is None on every "
            "member — usd_balance=$0 with healthy probes is a free-tier "
            "signal, not a real outage.",
        )

    def test_all_in_window_exhausted_fires_banner(self) -> None:
        """H3 — every chain member has credits<=low_water AND a future reset
        AND is consequently NOT available_for_calls. Banner SHOULD fire."""
        m1 = "groq:llama-3.3-70b-versatile"
        m2 = "groq:meta-llama/llama-4-scout-17b-16e-instruct"
        future = time.monotonic() + 3600.0
        h1 = _healthy_now(m1)
        h1.latency_ms = 100
        h1.credits_remaining = 100.0
        h1.credits_unit = "tokens_day"
        h1.credits_low_water = GROQ_TOKENS_LOW_WATER
        h1.credits_reset_at = future
        h1.credits_observed_at = time.monotonic()
        h2 = _healthy_now(m2)
        h2.latency_ms = 110
        h2.credits_remaining = 50.0
        h2.credits_unit = "tokens_day"
        h2.credits_low_water = GROQ_TOKENS_LOW_WATER
        h2.credits_reset_at = future
        h2.credits_observed_at = time.monotonic()
        block = self._build([m1, m2], {m1: h1, m2: h2})
        self.assertTrue(
            block["chain_credit_exhausted"],
            "Banner SHOULD fire when every chain member is strictly "
            "credit-exhausted with an in-window future reset.",
        )

    def test_chain_block_carries_current_primary_and_last_probe(self) -> None:
        """KI-122 — verify the new wire fields are populated."""
        nim = "qwen/qwen3-next-80b-a3b-instruct"
        nh = _healthy_now(nim)
        nh.latency_ms = 250
        block = self._build([nim], {nim: nh})
        self.assertEqual(block["current_primary"], nim)
        self.assertTrue(block["current_primary_available"])
        self.assertIsNotNone(block["last_probe_at"])
        self.assertIsNotNone(block["last_probe_age_seconds"])
        self.assertGreaterEqual(block["last_probe_age_seconds"], 0.0)


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