tests/test_server_functions.py

"""The three language server functions behind the gr.HTML bridge, and the tokens on the directive.

Plan 02-06. The bridge rules from 01-08 hold for every function here: ONE positional payload
(a dict from JS, a list for several arguments, ``[]`` for none), and ``{"error": ...}`` on every
failure because a raise becomes ``undefined`` in the browser. Two rules are this plan's own:

* ``analyze``, ``translate``, ``language_info`` and the analyze stage of ``turn`` all call
  ``warm_language()`` first, so a request that lands while ``Blocks.load`` is still filling the
  caches waits on the lock instead of racing ``lru_cache`` misses into a second JMdict load.
* ``language_info`` never reports a half-filled cache, for the same reason.

Quick-loop hygiene: the real ``warm_language()`` loads the translator (+100 MB), so the tests
that only need the analyzer run with ``blocks.warm_language`` redirected to a dictionary-only
stub (``dictionary_only_warm``) and the ones that exercise the real warm-up carry the ``mt``
marker. ``test_server_functions_wait_for_warm_language`` is what proves the real one is called.
"""

from __future__ import annotations

import sys
import threading
import time
from types import SimpleNamespace

import pytest

from japanese_avatar.nlp import warm
from japanese_avatar.ui import blocks
from japanese_avatar.ui.blocks import MAX_TEXT_CHARS, analyze, language_info, translate, turn

EAT = "食べました"
STATION = "駅はどこですか。"
LANGUAGE_INFO_KEYS = {
    "tokenizer_dict_bytes",
    "jmdict_entries",
    "jmdict_bytes",
    "mt_model_bytes",
    "warm",
    "rss_mb",
    "cpu_cores",
    "memory",
    "pins",
}
WARM_KEYS = {
    "tokenizer_s",
    "jmdict_s",
    "translator_s",
    "total_s",
    "rss_before_mb",
    "rss_mb",
    "rss_delta_mb",
    "expected_rss_delta_mb",
    "cpu_cores",
    "memory",
}
JMDICT_ENTRIES = 218672
BAD_PAYLOADS = ["", "   ", None, [], {"text": "あ" * (MAX_TEXT_CHARS + 1)}]


@pytest.fixture
def dictionary_only_warm(monkeypatch, analyzer):
    """Redirect ``blocks.warm_language`` to a stub so the quick loop never loads the translator.

    The session ``analyzer`` fixture has already built the tokenizer and the compact JMdict, so
    the analysis path is real; only the translator load is skipped. Returns the call log.
    """
    calls: list[float] = []

    def stub() -> dict:
        calls.append(time.perf_counter())
        return {"total_s": 0.0, "stub": True}

    monkeypatch.setattr(blocks, "warm_language", stub)
    return calls


def _one_tappable(tokens: list[dict]) -> dict:
    tappable = [t for t in tokens if t["tappable"]]
    assert len(tappable) == 1, tokens
    return tappable[0]


# --------------------------------------------------------------------------------- analyze


def test_analyze_accepts_the_bridge_payload_shapes(dictionary_only_warm):
    for payload in ({"text": EAT}, [EAT], EAT):
        out = analyze(payload)
        assert set(out) == {"tokens", "timings"}, out
        assert len(out["tokens"]) == 1
        assert _one_tappable(out["tokens"])["lemma"] == "食べる"
        assert {"analyze_ms", "server_total_ms"} <= set(out["timings"])
        assert out["timings"]["analyze_ms"] >= 0
    assert len(dictionary_only_warm) == 3, "analyze must wait for warm_language() on every call"


def test_analyze_rejects_bad_input_with_a_structured_error(dictionary_only_warm):
    for payload in BAD_PAYLOADS:
        out = analyze(payload)
        assert "error" in out and "tokens" not in out, (payload, out)
    assert "characters" in analyze({"text": "あ" * (MAX_TEXT_CHARS + 1)})["error"]
    assert dictionary_only_warm == [], "a rejected payload must not touch the warm-up"


# ------------------------------------------------------------------------------- translate


def test_translate_rejects_bad_input_with_a_structured_error():
    loaded_before = "ctranslate2" in sys.modules
    for payload in BAD_PAYLOADS:
        out = translate(payload)
        assert "error" in out and "text" not in out, (payload, out)
    assert "characters" in translate({"text": "あ" * (MAX_TEXT_CHARS + 1)})["error"]
    if not loaded_before:
        assert "ctranslate2" not in sys.modules, "rejecting bad input must not load the model"


@pytest.mark.mt
def test_translate_happy_path():
    out = translate({"text": STATION, "line_id": "L3"})
    assert set(out) == {"text", "line_id", "timings"}, out
    assert "station" in out["text"].lower(), out["text"]
    assert out["line_id"] == "L3"
    assert out["timings"]["translate_ms"] > 0
    as_list = translate([STATION, "L4"])
    assert as_list["line_id"] == "L4" and "station" in as_list["text"].lower()
    plain = translate(STATION)
    assert plain["line_id"] is None and "station" in plain["text"].lower()
    print(f"\ntranslate {STATION!r} -> {out['text']!r} in {out['timings']['translate_ms']:.1f} ms")


# --------------------------------------------------------------------------- language_info


@pytest.mark.mt
def test_language_info_shape():
    for payload in ([], None):
        info = language_info(payload) if payload is not None else language_info()
        assert set(info) == LANGUAGE_INFO_KEYS, sorted(info)
        assert info["jmdict_entries"] == JMDICT_ENTRIES
        assert info["mt_model_bytes"] > 70_000_000
        assert info["tokenizer_dict_bytes"] > 100_000_000, "system.dic is ~217 MB"
        assert info["jmdict_bytes"] > 1_000_000
        assert info["pins"]["sudachidict_core"] == "20260723"
        assert set(info["pins"]) == {
            "sudachipy",
            "sudachidict_core",
            "jmdict",
            "jlpt_vocab",
            "opus_mt_revision",
        }
        assert isinstance(info["warm"]["total_s"], int | float)
        assert info["warm"]["expected_rss_delta_mb"] == 410
    print(f"\nlanguage_info: {info}")


def test_language_info_shape_without_warm(monkeypatch, analyzer):
    """The quick-loop half: same exact key set, the warm dict is whatever warm_language returns."""
    loaded_before = "ctranslate2" in sys.modules
    stub = {"total_s": 0.0, "stub": True}
    monkeypatch.setattr(blocks, "warm_language", lambda: stub)
    info = language_info([])
    assert set(info) == LANGUAGE_INFO_KEYS, sorted(info)
    assert info["warm"] is stub
    assert info["jmdict_entries"] == JMDICT_ENTRIES
    assert info["mt_model_bytes"] > 70_000_000, "a byte count from disk, not a model load"
    assert info["tokenizer_dict_bytes"] > 100_000_000
    assert info["pins"]["sudachidict_core"] == "20260723"
    assert info["pins"]["jlpt_vocab"] == "2025.08.01.0"
    assert len(info["pins"]["opus_mt_revision"]) == 40
    if not loaded_before:
        assert "ctranslate2" not in sys.modules


def test_language_info_never_raises(monkeypatch):
    def boom() -> dict:
        raise RuntimeError("warm-up exploded")

    monkeypatch.setattr(blocks, "warm_language", boom)
    out = language_info([])
    assert set(out) == {"error"} and "warm-up exploded" in out["error"]


# ------------------------------------------------------------------------------- warm-up


@pytest.mark.mt
def test_warm_language_is_idempotent_and_measured():
    first = warm.warm_language()
    assert set(first) >= WARM_KEYS, sorted(first)
    assert first["expected_rss_delta_mb"] == 410
    assert isinstance(first["total_s"], int | float)
    assert warm.warm_language() is first, "the second call must return the same object"
    if first["rss_mb"] is None:
        print("\nwarm_language: rss unreadable on this machine (no /proc, no psutil)")
    else:
        print(
            f"\nwarm_language: rss {first['rss_before_mb']:.0f} -> {first['rss_mb']:.0f} MB, "
            f"delta {first['rss_delta_mb']:.0f} MB (research expected +~410 = 62 mmap + "
            f"241 JMdict + 103 CT2); tokenizer {first['tokenizer_s']:.3f}s "
            f"jmdict {first['jmdict_s']:.3f}s translator {first['translator_s']:.3f}s"
        )


def test_warm_language_first_call_is_serialised(monkeypatch):
    """Two concurrent FIRST callers: each component warms exactly once (the lock, not lru_cache)."""
    counts = {"tokenizer": 0, "jmdict": 0, "translator": 0}
    tally_lock = threading.Lock()

    def slow(name: str):
        def component() -> float:
            with tally_lock:
                counts[name] += 1
            time.sleep(0.05)
            return 0.05

        return component

    monkeypatch.setattr(warm, "_warm_tokenizer", slow("tokenizer"))
    monkeypatch.setattr(warm, "_warm_jmdict", slow("jmdict"))
    monkeypatch.setattr(warm, "_warm_translator", slow("translator"))
    warm._warm_once.cache_clear()
    try:
        results: list[dict] = []
        threads = [
            threading.Thread(target=lambda: results.append(warm.warm_language())) for _ in range(2)
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=10)
        assert len(results) == 2
        assert results[0] is results[1]
        assert counts == {"tokenizer": 1, "jmdict": 1, "translator": 1}, counts
        assert results[0]["errors"] == []
    finally:
        warm._warm_once.cache_clear()  # never leave the stub result where language_info reads it


def test_warm_language_survives_a_failing_component(monkeypatch):
    """warm_synthesizer's rule: start-up must not take the page down."""

    def broken() -> float:
        raise RuntimeError("no model here")

    monkeypatch.setattr(warm, "_warm_tokenizer", lambda: 0.0)
    monkeypatch.setattr(warm, "_warm_jmdict", lambda: 0.0)
    monkeypatch.setattr(warm, "_warm_translator", broken)
    warm._warm_once.cache_clear()
    try:
        out = warm.warm_language()
        assert out["translator_s"] is None
        assert out["errors"] == ["translator: no model here"]
        assert isinstance(out["total_s"], float)
    finally:
        warm._warm_once.cache_clear()


def test_rss_mb_is_a_number_or_none():
    value = warm.rss_mb()
    assert value is None or (isinstance(value, float) and value > 10.0), value
    print(f"\nrss_mb() = {value}")


# ------------------------------------------------------------------- tokens on the directive


def test_turn_directive_carries_tokens(dictionary_only_warm):
    pytest.importorskip("voicevox_core")
    directive = turn(EAT)
    assert "error" not in directive, directive
    assert len(directive["tokens"]) == 1
    assert _one_tappable(directive["tokens"])["lemma"] == "食べる"
    assert directive["timings"]["analyze_ms"] >= 0
    assert len(dictionary_only_warm) == 1


def test_turn_speaks_when_analysis_fails(monkeypatch, dictionary_only_warm):
    pytest.importorskip("voicevox_core")

    def broken(_text: str) -> list[dict]:
        raise RuntimeError("analyzer down")

    monkeypatch.setattr(blocks.nlp_analyzer, "analyze", broken)
    directive = turn(EAT)
    assert "error" not in directive, directive
    assert directive["tokens"] == []
    assert directive["audio_url"].startswith("data:audio/wav;base64,")


def test_server_functions_wait_for_warm_language(monkeypatch, analyzer):
    """analyze, translate and turn each call warm_language() before touching a singleton."""
    counter = {"n": 0}

    def counting() -> dict:
        counter["n"] += 1
        return {"total_s": 0.0, "stub": True}

    monkeypatch.setattr(blocks, "warm_language", counting)
    monkeypatch.setattr(blocks.nlp_translate, "translate", lambda text: "Hello.")
    monkeypatch.setattr(
        blocks,
        "synthesize",
        lambda text, speed, timings: SimpleNamespace(wav_bytes=b"RIFF", audio_query={}),
    )
    monkeypatch.setattr(blocks, "build_timeline", lambda query: [])
    monkeypatch.setattr(blocks, "timeline_to_dicts", lambda timeline: [])

    assert "tokens" in analyze({"text": "こんにちは"})
    assert translate({"text": "こんにちは"})["text"] == "Hello."
    directive = turn({"text": "こんにちは"})
    assert "error" not in directive, directive
    assert directive["tokens"] and directive["tokens"][0]["surface"] == "こんにちは"
    assert counter["n"] == 3, counter