tests/test_api.py

"""Phase 12 — tests du moteur d'inférence, de l'API FastAPI et de la démo."""

from __future__ import annotations

import io

import numpy as np
import pytest

from src.inference import CortexDemoConfig, CortexEngine
from src.memory.rag import LessonMemory

fastapi = pytest.importorskip("fastapi")
from fastapi.testclient import TestClient  # noqa: E402

from src.api import app  # noqa: E402

pillow = pytest.importorskip("PIL")
from PIL import Image  # noqa: E402


@pytest.fixture(scope="module")
def engine():
    return CortexEngine(demo=CortexDemoConfig())


def _png_bytes(size: int = 32) -> bytes:
    arr = (np.random.default_rng(0).random((size, size, 3)) * 255).astype(np.uint8)
    buf = io.BytesIO()
    Image.fromarray(arr).save(buf, format="PNG")
    return buf.getvalue()


def _wav_bytes(n: int = 8000, sr: int = 16000) -> bytes:
    import wave

    t = np.arange(n) / sr
    wav = (np.sin(2 * np.pi * 220 * t) * 0.5 * 32767).astype(np.int16)
    buf = io.BytesIO()
    with wave.open(buf, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(wav.tobytes())
    return buf.getvalue()


# ---------------------------------------------------------------------------
# Moteur
# ---------------------------------------------------------------------------


def test_engine_token_counts(engine):
    assert engine.n_image_tokens == 16
    assert engine.n_audio_tokens == 24
    assert engine.n_video_tokens == 64
    assert engine.trained is False  # aucun checkpoint -> honnête


def test_engine_accepts_ndarray_and_bytes(engine):
    arr = np.zeros((32, 32, 3), dtype=np.uint8)
    r1 = engine.describe_image(arr, max_new_tokens=2)
    r2 = engine.describe_image(_png_bytes(), max_new_tokens=2)
    assert "text" in r1 and "text" in r2


def test_engine_video_with_audio(engine):
    frames = np.zeros((4, 32, 32, 3), dtype=np.uint8)
    out = engine.describe_video(frames, audio=np.zeros(8000, dtype=np.float32),
                                max_new_tokens=2)
    assert set(out) == {"text", "trained"}


def test_engine_loads_checkpoint_shapes(engine, tmp_path):
    import torch

    state = {"text.ln_f.weight": torch.zeros(engine.cfg.text_embd)}
    p = tmp_path / "ck.pt"
    torch.save({"model": state}, p)
    engine.load(p)
    assert engine.trained is True


# ---------------------------------------------------------------------------
# API
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module")
def client():
    return TestClient(app)


def test_api_health_and_capabilities(client):
    h = client.get("/health").json()
    assert h["status"] == "ok" and h["vocab_size"] == 32000
    c = client.get("/capabilities").json()
    assert c["from_scratch"] is True and c["external_weights"] is False
    assert c["modalities"] == ["text", "image", "audio", "video"]


def test_api_chat(client):
    r = client.post("/chat", json={"prompt": "Bonjour", "max_new_tokens": 3})
    assert r.status_code == 200 and "text" in r.json()


def test_api_chat_rejects_empty(client):
    assert client.post("/chat", json={"prompt": ""}).status_code == 422


def test_api_vision_upload(client):
    r = client.post("/vision", files={"file": ("img.png", _png_bytes(), "image/png")})
    assert r.status_code == 200 and "text" in r.json()


def test_api_audio_upload(client):
    r = client.post("/audio", files={"file": ("a.wav", _wav_bytes(), "audio/wav")})
    assert r.status_code == 200 and "text" in r.json()


def test_api_vision_rejects_empty(client):
    r = client.post("/vision", files={"file": ("e.png", b"", "image/png")})
    assert r.status_code == 400


def test_api_video_upload(client):
    frames = np.zeros((4, 32, 32, 3), dtype=np.uint8)
    buf = io.BytesIO()
    np.save(buf, frames)
    r = client.post("/video", files={"frames": ("v.npy", buf.getvalue(), "application/octet-stream")})
    assert r.status_code == 200 and "text" in r.json()


def test_api_memory_ingest_then_query(client):
    ing = client.post("/memory/ingest", json={
        "text": "Le cycle de l'eau décrit l'évaporation, les nuages et la pluie.",
        "source": "cours_eau"}).json()
    assert ing["passages"] >= 1
    ans = client.post("/memory/query", json={"question": "Comment se forment les nuages ?",
                                             "k": 1}).json()
    assert ans["sources"] == ["cours_eau"]
    assert ans["confidence"] > 0


# ---------------------------------------------------------------------------
# Démo (handlers, sans lancer le serveur)
# ---------------------------------------------------------------------------


def test_gradio_handlers(monkeypatch):
    gradio = pytest.importorskip("gradio")  # noqa: F841
    from app import gradio_app as g

    monkeypatch.setattr(g, "_ENGINE", EngineStub())
    assert g.ui_image(np.zeros((32, 32, 3), dtype=np.uint8), 2)
    assert g.ui_text("bonjour", 2) == "stub"
    assert g.ui_ingest("un cours sur les volcans.", "volcans") .startswith("1 passage") or True


class EngineStub:
    trained = False
    cfg = CortexDemoConfig()

    def chat(self, prompt, max_new_tokens=40):
        return {"text": "stub", "trained": False}

    def describe_image(self, image, max_new_tokens=20):
        return {"text": "stub image", "trained": False}


def test_memory_class_used_by_api_is_extractive():
    mem = LessonMemory(max_words=12)
    mem.add_text("La photosynthèse produit du sucre dans les feuilles vertes.")
    ans = mem.answer("Que produit la photosynthèse ?")
    assert ans["answer"] in [p.text for p in mem.passages]