tests/test_engine.py

"""The model calls, with Inference Providers replaced by a fake: backup, budget, and billing."""
from __future__ import annotations

import io
import json
import sys
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pytest
import soundfile as sf

import contract
import engine
import judges

PANEL = list(judges.DEFAULT_PANEL)
REVIEW = json.dumps({**{seat: {"quote": 0, "strength": "Real.", "concern": "Untested.",
                               "improvement": "Test it.", "score": 7} for seat in contract.SEATS},
                     **{field: "A line." for field, _ in contract.DISCUSSION}})
PRIMARY, BACKUP, LAST = (f"{route.provider}/{route.model}" for route in engine.ROUTES)


def chat_reply(content=REVIEW):
    return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content=content))])


class FakeClient:
    """Stands in for InferenceClient. `script` maps a model to a reply, or to an exception to raise."""

    def __init__(self, calls, script, provider=None):
        self.calls, self.script, self.provider = calls, script, provider

    def chat_completion(self, messages, model, **options):
        route = f"{self.provider}/{model}"
        self.calls.append({"model": route, "messages": messages, **options})
        outcome = self.script.get(route, chat_reply())
        if isinstance(outcome, Exception):
            raise outcome
        return outcome

    def text_to_speech(self, text, model, extra_body):
        self.calls.append({"model": model, "text": text, "provider": self.provider, **extra_body})
        if isinstance(self.script.get(self.provider), Exception):
            raise self.script[self.provider]
        buffer = io.BytesIO()
        sf.write(buffer, np.ones(2400, dtype="float32") * 0.1, 24000, format="WAV")
        return buffer.getvalue()

    def automatic_speech_recognition(self, audio, model):
        self.calls.append({"model": model, "audio": audio})
        outcome = self.script.get(model, SimpleNamespace(text=f" Hello from {model}. "))
        if isinstance(outcome, Exception):
            raise outcome
        return outcome


@pytest.fixture
def providers(monkeypatch, tmp_path):
    calls, script = [], {}
    monkeypatch.setattr(engine, "_client", lambda provider, timeout=None: FakeClient(calls, script, provider))
    monkeypatch.setattr(engine.config, "HF_TOKEN", "hf_test")
    monkeypatch.setattr(engine, "BUDGET", engine.Budget(tmp_path / "spend.json", 1.80))
    return calls, script


def spent(tmp_path):
    return json.loads((tmp_path / "spend.json").read_text())["usd"]


# ---------------------------------------------------------------- the review


def test_the_primary_model_writes_the_review(providers):
    calls, _ = providers
    assert engine.review(judges.EXAMPLE_PITCH, PANEL).seats[0].score == 7
    assert [c["model"] for c in calls] == [PRIMARY]


@pytest.mark.parametrize("failure", [RuntimeError("429 rate limited"), TimeoutError(), chat_reply("not json")])
def test_the_backup_answers_when_the_primary_cannot(providers, failure):
    calls, script = providers
    script[PRIMARY] = failure
    assert engine.review(judges.EXAMPLE_PITCH, PANEL).seats[0].score == 7
    assert [c["model"] for c in calls] == [PRIMARY, BACKUP]


def test_when_nobody_can_answer_the_visitor_is_told_kindly(providers):
    _, script = providers
    script[PRIMARY] = script[BACKUP] = script[LAST] = RuntimeError("busy")
    with pytest.raises(ValueError, match="swamped"):
        engine.review(judges.EXAMPLE_PITCH, PANEL)


def test_thinking_is_kept_short_on_every_route():
    """Left on, it spends the whole token budget thinking and never writes the review."""
    for route in engine.ROUTES:
        options = route.options
        assert options.get("chat_template_kwargs", {}).get("enable_thinking") is False \
            or options.get("reasoning_effort") == "low"


def test_glm_leads_and_every_route_asks_for_json():
    assert engine.ROUTES[0].model == "zai-org/GLM-5.3-Flash"
    formats = [route.options["response_format"]["type"] for route in engine.ROUTES]
    assert all(f in ("json_object", "json_schema") for f in formats)
    assert all(route.schema_in_prompt for route, f in zip(engine.ROUTES, formats) if f == "json_object")


def test_the_pitch_is_sent_as_data_not_instructions(providers):
    calls, _ = providers
    engine.review(judges.EXAMPLE_PITCH, PANEL)
    system, user = calls[0]["messages"]
    assert judges.EXAMPLE_PITCH not in system["content"]
    assert json.loads(user["content"])["pitch"] == judges.EXAMPLE_PITCH


# ---------------------------------------------------------------- speech


def test_each_judge_is_voiced_in_their_own_voice_with_pauses(providers):
    calls, _ = providers
    rate, audio, starts = engine.speak([("investor", "One."), ("builder", "Two.")])
    assert [c["voice"] for c in calls] == [judges.JUDGES["investor"].voice, judges.JUDGES["builder"].voice]
    assert rate == 24000 and len(audio) == 2 * (2400 + int(engine.config.GAP_SECONDS * 24000))
    assert starts == [0.0, round(0.1 + engine.config.GAP_SECONDS, 2)], "the second line starts after the first and its pause"


def test_voices_fall_back_to_the_backup_provider_with_its_own_voices(providers):
    calls, script = providers
    script["deepinfra"] = TimeoutError()
    engine.speak([("investor", "One."), ("customer", "Two.")])
    backup = [c for c in calls if c["provider"] == "fal-ai"]
    assert [c["voice"] for c in backup] == [judges.JUDGES["investor"].backup_voice, judges.JUDGES["customer"].backup_voice]


def test_first_prefers_the_first_attempt_when_it_is_healthy():
    assert engine._first([lambda: "primary", lambda: "backup"], head_start=1) == "primary"


def test_first_starts_the_backup_at_once_when_the_primary_fails():
    def broken():
        raise TimeoutError()
    started = engine.time.perf_counter()
    assert engine._first([broken, lambda: "backup"], head_start=30) == "backup"
    assert engine.time.perf_counter() - started < 1, "a failure must not wait out the head start"


def test_first_lets_the_backup_overtake_a_stalled_primary():
    def stalled():
        engine.time.sleep(3)
        return "primary"
    started = engine.time.perf_counter()
    assert engine._first([stalled, lambda: "backup"], head_start=0.2) == "backup"
    assert engine.time.perf_counter() - started < 1


def test_first_raises_when_everything_fails():
    def broken():
        raise TimeoutError()
    with pytest.raises(TimeoutError):
        engine._first([broken, broken], head_start=0.1)


def test_every_judge_has_their_own_voice_and_a_backup_voice_the_backup_provider_offers():
    assert len({j.voice for j in judges.JUDGES.values()}) == len(judges.JUDGES)
    american = {j.backup_voice for j in judges.JUDGES.values()}
    assert len(american) == len(judges.JUDGES) and all(v[:3] in ("af_", "am_") for v in american)


def test_transcription_uses_whisper_and_leaves_the_backup_idle_when_it_is_quick(providers):
    calls, _ = providers
    assert engine.transcribe(np.zeros(16000, dtype="float32")) == "Hello from openai/whisper-large-v3-turbo."
    assert [c["model"] for c in calls] == ["openai/whisper-large-v3-turbo"], "no second bill when the first is quick"
    assert str(calls[0]["audio"]).endswith(".flac")


def test_a_long_recording_gives_whisper_a_longer_head_start(providers, monkeypatch):
    head_starts = []
    monkeypatch.setattr(engine, "_first", lambda attempts, head_start: head_starts.append(head_start) or "Hi")
    engine.transcribe(np.zeros(16000 * 60 * 4, dtype="float32"))
    assert head_starts == [engine.config.ASR_HEAD_START + 4 * engine.config.ASR_SECONDS_PER_MINUTE]


def test_transcription_survives_whisper_failing(providers):
    _, script = providers
    script["openai/whisper-large-v3-turbo"] = TimeoutError()
    assert engine.transcribe(np.zeros(16000, dtype="float32")) == "Hello from Qwen/Qwen3-ASR-1.7B."


def test_transcription_failing_everywhere_is_explained(providers):
    _, script = providers
    for _, model in engine.ASR:
        script[model] = TimeoutError()
    with pytest.raises(ValueError, match="couldn't transcribe"):
        engine.transcribe(np.zeros(16000, dtype="float32"))


# ---------------------------------------------------------------- spending


def test_voices_never_run_the_backup_alongside_a_healthy_provider(providers):
    calls, _ = providers
    engine.speak([("investor", "One."), ("builder", "Two.")])
    assert {c["provider"] for c in calls} == {"deepinfra"}


def test_audio_is_billed_by_the_second_even_when_it_fails(providers, tmp_path, monkeypatch):
    clock = iter([100.0, 110.0])                                       # a request that hangs for ten seconds
    monkeypatch.setattr(engine.time, "perf_counter", lambda: next(clock))
    with pytest.raises(TimeoutError):
        engine._billed("deepinfra", lambda: (_ for _ in ()).throw(TimeoutError()))
    assert spent(tmp_path) == pytest.approx(10 * engine.PER_SECOND["deepinfra"])


def test_every_review_attempt_is_counted_at_its_route_price(providers, tmp_path):
    _, script = providers
    engine.review(judges.EXAMPLE_PITCH, PANEL)
    assert spent(tmp_path) == pytest.approx(engine.ROUTES[0].cost)
    script[PRIMARY] = chat_reply("")                                   # paid for, but unreadable
    engine.review(judges.EXAMPLE_PITCH + " Again.", PANEL)
    assert spent(tmp_path) == pytest.approx(2 * engine.ROUTES[0].cost + engine.ROUTES[1].cost)


def test_the_den_stops_at_the_budget(providers, monkeypatch, tmp_path):
    calls, _ = providers
    monkeypatch.setattr(engine, "BUDGET", engine.Budget(tmp_path / "spend.json", 0.015))
    engine.review(judges.EXAMPLE_PITCH, PANEL)
    engine.review(judges.EXAMPLE_PITCH, PANEL)                        # two cents: over the limit
    for call in (lambda: engine.review(judges.EXAMPLE_PITCH, PANEL), lambda: engine.speak([("investor", "Hi.")]),
                 lambda: engine.transcribe(np.zeros(160, dtype="float32"))):
        with pytest.raises(ValueError, match="budget"):
            call()
    assert len(calls) == 2, "nothing is sent once the budget is spent"


def test_an_account_out_of_credit_is_explained_and_unbilled(providers, tmp_path):
    _, script = providers
    refused = RuntimeError("402")
    refused.response = SimpleNamespace(status_code=402)
    for route in (PRIMARY, BACKUP, LAST):
        script[route] = refused
    with pytest.raises(ValueError, match="budget"):
        engine.review(judges.EXAMPLE_PITCH, PANEL)
    assert not (tmp_path / "spend.json").exists(), "a refused request costs nothing"


def test_spending_survives_a_restart_and_resets_each_month(providers, tmp_path):
    path = tmp_path / "spend.json"
    engine.Budget(path, 1.0).add(0.4)
    assert engine.Budget(path, 1.0).add(0.1) == pytest.approx(0.5)
    path.write_text(json.dumps({"month": "1999-01", "usd": 5}))
    engine.Budget(path, 1.0).check()


def test_no_token_is_a_friendly_message_and_sends_nothing(providers, monkeypatch):
    calls, _ = providers
    monkeypatch.setattr(engine.config, "HF_TOKEN", "")
    with pytest.raises(ValueError, match="isn't connected"):
        engine.review(judges.EXAMPLE_PITCH, PANEL)
    assert calls == []