bikoret/kushya_generator/test/test_dataset.py

"""
Validate the kushya training dataset structure.

Architecture §3.6 — every training pair must satisfy the four requirements
or be rejected. These tests run prepare_training.py's validator against
the seed pairs to catch dataset regressions.
"""

import json
import sys
from pathlib import Path

# Add parent data/ dir to path so we can import prepare_training
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "data"))

from prepare_training import (  # noqa: E402
    REQUIRED_TYPES,
    ValidationError,
    load_jsonl,
    to_training_format,
    validate_pair,
)


SEED_PATH = (
    Path(__file__).resolve().parent.parent / "data" / "training_pairs" / "seed_pairs.jsonl"
)


def test_seed_pairs_file_exists():
    """The seed pairs file must exist — it bootstraps the training pipeline."""
    assert SEED_PATH.exists(), f"seed_pairs.jsonl not found at {SEED_PATH}"


def test_seed_pairs_validate():
    """Every seed pair must pass validation. If a pair fails, the dataset
    has regressed and the training pipeline would silently degrade."""
    failures = []
    for line_no, pair in load_jsonl(SEED_PATH):
        try:
            validate_pair(pair, "seed_pairs.jsonl", line_no)
        except ValidationError as e:
            failures.append(str(e))
    assert not failures, "seed_pairs.jsonl has invalid records:\n" + "\n".join(failures)


def test_seed_pairs_cover_all_kushya_types():
    """The seed set should exercise every kushya type, so the training
    split is stratified across the categories."""
    types_seen = set()
    for _, pair in load_jsonl(SEED_PATH):
        types_seen.add(pair["type"])

    # We don't require ALL types to be in the small seed (10 pairs),
    # but we do require at least 3 types to ensure stratification works.
    assert len(types_seen) >= 3, (
        f"Only {len(types_seen)} types in seed: {types_seen}. "
        f"Need at least 3 for stratified split."
    )

    # Every type used must be a known type (validator already checks this,
    # but we want a clear failure if a typo creeps in).
    unknown = types_seen - REQUIRED_TYPES
    assert not unknown, f"Unknown types in seed: {unknown}"


def test_to_training_format_produces_expected_shape():
    """The training format must follow Architecture §3.6:
      input:  'Proposed interpretation: ...'
      output: 'Kushya: ...'
    """
    sample = {
        "interpretation": "Cooking on Shabbat is forbidden.",
        "kushya": "What about cooking by direct sunlight where no kindling occurs in Shemot 35:3?",
        "torah_refs": ["Shemot 35:3"],
        "type": "scope_narrowing",
    }
    formatted = to_training_format(sample)
    assert formatted["input"].startswith("Proposed interpretation:")
    assert formatted["output"].startswith("Kushya:")
    assert formatted["metadata"]["type"] == "scope_narrowing"
    assert formatted["metadata"]["torah_refs"] == ["Shemot 35:3"]


# ============================================================
# Validator regression tests — these should REJECT bad pairs.
# ============================================================

def _expect_rejection(pair, must_contain):
    try:
        validate_pair(pair, "test", 1)
    except ValidationError as e:
        assert must_contain in str(e), (
            f"Expected error containing {must_contain!r}, got: {e}"
        )
    else:
        raise AssertionError(f"Pair should have been rejected: {pair}")


def test_validator_rejects_empty_kushya():
    pair = {
        "interpretation": "Some interp.",
        "kushya": "",
        "torah_refs": ["Shemot 21:1"],
        "type": "internal_tension",
    }
    _expect_rejection(pair, "kushya must be a non-empty string")


def test_validator_rejects_kushya_without_question_form():
    pair = {
        "interpretation": "Some interp.",
        "kushya": "There is a problem with this interpretation that needs clarification.",
        "torah_refs": ["Shemot 21:1"],
        "type": "internal_tension",
    }
    _expect_rejection(pair, "expressible as a question")


def test_validator_rejects_too_short_kushya():
    pair = {
        "interpretation": "Some interp.",
        "kushya": "Why?",
        "torah_refs": ["Shemot 21:1"],
        "type": "internal_tension",
    }
    _expect_rejection(pair, "too short")


def test_validator_rejects_no_torah_reference():
    pair = {
        "interpretation": "Some interp.",
        "kushya": "What does this even mean in context — there is internal tension here is there not?",
        "torah_refs": [],
        "type": "internal_tension",
    }
    _expect_rejection(pair, "must cite a specific Torah passage")


def test_validator_rejects_unknown_type():
    pair = {
        "interpretation": "Some interp.",
        "kushya": "What does this mean in Shemot 21:1, given the apparent tension?",
        "torah_refs": ["Shemot 21:1"],
        "type": "made_up_type",
    }
    _expect_rejection(pair, "not in known types")