rpa_vision_v3/tests/unit/test_role_mapper.py

"""Tests du role_mapper : reconstruction de champs ANCRÉS sur l'OCR.

Principe cardinal (cf. gate vert 30/06) : le VLM ne fournit QUE des ids de tokens OCR
(value_ids) ; la valeur est reconstruite côté Python depuis l'OCR. Aucun texte produit
par le VLM ne doit pouvoir entrer dans une valeur -> 0 hallucination par construction.
"""
import pytest

from core.extraction.role_mapper import (
    OcrToken,
    build_role_prompt,
    map_roles,
    reconstruct_fields,
    tokens_from_grid,
)


def _tok(tid, text, conf=0.9, bbox=(0, 0, 10, 10)):
    return OcrToken(id=tid, text=text, confidence=conf, bbox=bbox)


def test_reconstruit_value_concatene_tokens_dans_lordre():
    tokens = [_tok(0, "DUPONT"), _tok(1, "Jean")]
    fields = reconstruct_fields(tokens, [{"label": "Nom complet", "value_ids": [0, 1]}])
    assert len(fields) == 1
    assert fields[0].label == "Nom complet"
    assert fields[0].value == "DUPONT Jean"
    assert fields[0].anchored is True


def test_ignore_les_ids_hors_plage_et_les_liste():
    tokens = [_tok(0, "DUPONT")]
    fields = reconstruct_fields(tokens, [{"label": "Nom", "value_ids": [0, 99]}])
    assert fields[0].value == "DUPONT"
    assert fields[0].invalid_ids == [99]
    assert fields[0].anchored is True


def test_value_ids_vide_donne_champ_non_ancre():
    tokens = [_tok(0, "DUPONT")]
    fields = reconstruct_fields(tokens, [{"label": "Poids", "value_ids": []}])
    assert fields[0].value == ""
    assert fields[0].anchored is False


def test_aucun_id_valide_donne_champ_non_ancre():
    tokens = [_tok(0, "DUPONT")]
    fields = reconstruct_fields(tokens, [{"label": "Poids", "value_ids": [7, 8]}])
    assert fields[0].anchored is False
    assert fields[0].value == ""
    assert fields[0].invalid_ids == [7, 8]


def test_dedup_ids_en_preservant_lordre():
    tokens = [_tok(0, "DUPONT"), _tok(1, "Jean")]
    fields = reconstruct_fields(tokens, [{"label": "X", "value_ids": [1, 1, 0]}])
    assert fields[0].value == "Jean DUPONT"
    assert fields[0].value_ids == [1, 0]


def test_confidence_est_le_min_des_tokens_ancres():
    tokens = [_tok(0, "A", conf=0.95), _tok(1, "B", conf=0.70)]
    fields = reconstruct_fields(tokens, [{"label": "X", "value_ids": [0, 1]}])
    assert fields[0].confidence == pytest.approx(0.70)


def test_bbox_englobante_des_tokens_ancres():
    tokens = [_tok(0, "A", bbox=(0, 0, 10, 10)), _tok(1, "B", bbox=(20, 5, 40, 15))]
    fields = reconstruct_fields(tokens, [{"label": "X", "value_ids": [0, 1]}])
    assert fields[0].bbox == (0, 0, 40, 15)


def test_invariant_aucun_texte_hors_ocr():
    # 'value' fournie par le VLM est ignorée : seul value_ids compte.
    tokens = [_tok(0, "DUPONT")]
    fields = reconstruct_fields(
        tokens, [{"label": "Nom", "value_ids": [0], "value": "HALLUCINATION"}]
    )
    assert fields[0].value == "DUPONT"


def test_tokens_from_grid_indexe_et_normalise_bbox():
    # grille extract_grid_from_image : bbox = 4 points EasyOCR
    grid = [
        [
            {"text": "Nom", "bbox": [[0, 0], [10, 0], [10, 8], [0, 8]],
             "confidence": 0.9, "row": 0, "col": 0},
            {"text": "DUPONT", "bbox": [[20, 0], [60, 0], [60, 8], [20, 8]],
             "confidence": 0.95, "row": 0, "col": 1},
        ],
    ]
    tokens = tokens_from_grid(grid)
    assert [t.id for t in tokens] == [0, 1]
    assert tokens[0].text == "Nom"
    assert tokens[1].bbox == (20, 0, 60, 8)


# --- map_roles : orchestrateur (client VLM injectable, donc testable hors-ligne) ---

def _fake_client(response, capture=None):
    """Faux client VLM : enregistre éventuellement le prompt reçu, renvoie une réponse fixe."""
    def client(image_path, prompt):
        if capture is not None:
            capture["prompt"] = prompt
            capture["image_path"] = image_path
        return response
    return client


def test_map_roles_reconstruit_via_client_injecte():
    tokens = [_tok(0, "DUPONT"), _tok(1, "Jean")]
    client = _fake_client('{"champs":[{"label":"Nom complet","value_ids":[0,1]}]}')
    fields = map_roles("img.png", tokens, client)
    assert len(fields) == 1
    assert fields[0].label == "Nom complet"
    assert fields[0].value == "DUPONT Jean"


def test_map_roles_tolere_les_fences_json():
    tokens = [_tok(0, "DUPONT")]
    client = _fake_client('```json\n{"champs":[{"label":"Nom","value_ids":[0]}]}\n```')
    fields = map_roles("img.png", tokens, client)
    assert fields[0].value == "DUPONT"


def test_map_roles_json_invalide_retourne_liste_vide():
    # robustesse batch : une réponse VLM non-JSON ne doit pas crasher.
    tokens = [_tok(0, "DUPONT")]
    client = _fake_client("désolé, je n'ai pas compris")
    fields = map_roles("img.png", tokens, client)
    assert fields == []


def test_build_role_prompt_inclut_les_tokens_avec_ids():
    tokens = [_tok(0, "Poids"), _tok(1, "72")]
    prompt = build_role_prompt(tokens)
    assert "Poids" in prompt and "72" in prompt
    assert "value_ids" in prompt  # on demande bien des ids, pas du texte recopié


def test_build_role_prompt_guide_liste_les_roles_attendus():
    tokens = [_tok(0, "X")]
    prompt = build_role_prompt(tokens, roles=["Nom", "IPP", "Poids"])
    assert "Nom" in prompt and "IPP" in prompt and "Poids" in prompt


def test_map_roles_passe_les_roles_au_prompt():
    tokens = [_tok(0, "X")]
    cap = {}
    client = _fake_client('{"champs":[]}', capture=cap)
    map_roles("img.png", tokens, client, roles=["Diagnostic", "GEMSA"])
    assert "Diagnostic" in cap["prompt"] and "GEMSA" in cap["prompt"]