tests/test_pubmed_language_normalization.py

from __future__ import annotations

import json
from pathlib import Path

import pytest

from extraction.schemas import (
    ExtractClaimsRequest,
    normalize_pubmed_language,
)


ROOT = Path(__file__).resolve().parents[1]
EXAMPLE = (
    ROOT
    / "contracts"
    / "extract-claims.request.example.json"
)


@pytest.mark.parametrize(
    ("pubmed_code", "expected"),
    [
        ("eng", "en"),
        ("fre", "fr"),
        ("fra", "fr"),
        ("ger", "de"),
        ("deu", "de"),
        ("spa", "es"),
        ("ita", "it"),
        ("por", "pt"),
        ("dut", "nl"),
        ("nld", "nl"),
        ("chi", "zh"),
        ("zho", "zh"),
        ("jpn", "ja"),
        ("ara", "ar"),
        ("rus", "ru"),
    ],
)
def test_normalize_pubmed_iso639_2(
    pubmed_code: str,
    expected: str,
) -> None:
    assert (
        normalize_pubmed_language(pubmed_code)
        == expected
    )


def test_existing_iso639_1_is_preserved() -> None:
    assert normalize_pubmed_language("en") == "en"
    assert normalize_pubmed_language("fr") == "fr"
    assert normalize_pubmed_language("pt-BR") == "pt-BR"


def test_extract_request_accepts_real_pubmed_eng() -> None:
    payload = json.loads(
        EXAMPLE.read_text(encoding="utf-8")
    )

    payload["article"]["language"] = "eng"

    parsed = ExtractClaimsRequest.model_validate(
        payload
    )

    assert parsed.article.language == "en"


def test_unknown_three_letter_code_is_rejected() -> None:
    payload = json.loads(
        EXAMPLE.read_text(encoding="utf-8")
    )

    payload["article"]["language"] = "xyz"

    with pytest.raises(Exception):
        ExtractClaimsRequest.model_validate(payload)