tests/test_cefr_preprocessing.py

import pytest

from tutor.ml.cefr.preprocessing import (
    chunk_text,
    normalize_level,
    passages_from_record,
)


@pytest.mark.parametrize(
    ("raw", "expected"),
    [
        ("A1", "A1"),
        ("c2", "C2"),
        (" b1 ", "B1"),
        ("B1+", "B1"),  # X+ floors to X (ADR 0003)
        ("C2+", "C2"),
        ("B", None),  # bare macro level: ambiguous, dropped
        ("NA", None),
        ("EMPTY", None),
        ("unrated", None),
        ("Unassessable", None),
        ("", None),
        (None, None),
        ("D1", None),  # not a CEFR level
    ],
)
def test_normalize_level(raw: str | None, expected: str | None) -> None:
    assert normalize_level(raw) == expected


def _text(n_sentences: int, words_per_sentence: int = 12) -> str:
    sentence = " ".join(f"w{i}" for i in range(words_per_sentence - 1))
    return ". ".join(f"{sentence} end" for _ in range(n_sentences)) + "."


def test_chunk_text_short_text_passes_through() -> None:
    text = _text(3)
    assert chunk_text(text, target_words=200, max_words=300) == [text]


def test_chunk_text_respects_max_and_preserves_words() -> None:
    text = _text(80)  # ~960 words
    chunks = chunk_text(text, target_words=200, max_words=300)
    assert len(chunks) > 1
    assert all(len(chunk.split()) <= 300 for chunk in chunks)
    # No word lost, none duplicated, order preserved:
    assert " ".join(chunks).split() == text.split()


def test_chunk_text_merges_small_tail() -> None:
    # 210 words then a 12-word tail: the tail must merge into the previous chunk.
    text = _text(18, words_per_sentence=12)  # 216 words total, target 200
    chunks = chunk_text(text, target_words=200, max_words=300, min_tail_words=50)
    assert len(chunks) == 1 or all(len(c.split()) >= 50 for c in chunks)


def test_chunk_text_hard_splits_a_giant_sentence() -> None:
    no_punctuation = " ".join(f"w{i}" for i in range(700))
    chunks = chunk_text(no_punctuation, target_words=200, max_words=300)
    assert all(len(chunk.split()) <= 300 for chunk in chunks)
    assert " ".join(chunks).split() == no_punctuation.split()


def _record(**overrides):
    base = {
        "text": _text(60),
        "level_raw": "B2",
        "lang": "en",
        "corpus": "cambridge_exams_en",
        "doc_id": "cambridge_exams_en:0",
        "source_format": "document-level",
    }
    base.update(overrides)
    return base


def test_passages_from_record_chunks_documents() -> None:
    passages = passages_from_record(**_record())
    assert len(passages) > 1
    assert {p.doc_id for p in passages} == {"cambridge_exams_en:0"}
    assert {p.level for p in passages} == {"B2"}  # chunks inherit the doc label


def test_passages_from_record_sentence_rows_never_chunked() -> None:
    passages = passages_from_record(**_record(source_format="sentence-level"))
    assert len(passages) == 1


def test_passages_from_record_truncation_arm() -> None:
    passages = passages_from_record(**_record(chunking=False))
    assert len(passages) == 1


def test_passages_from_record_drops_bad_labels_and_empty_text() -> None:
    assert passages_from_record(**_record(level_raw="B")) == []
    assert passages_from_record(**_record(level_raw="unrated")) == []
    assert passages_from_record(**_record(text="   ")) == []
    assert passages_from_record(**_record(text=None)) == []