tests/test_dataset.py

"""Tests for the synthetic dataset builder and rule engine."""

import json
import random
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.rules import humanize_rules, protect, restore

DATA = Path(__file__).resolve().parent.parent / "data"
EN_PHRASES = {
    "it is important to note that",
    "it is worth noting that",
    "moreover",
    "furthermore",
    "in conclusion",
    "leverage",
    "seamless",
    "robust",
    "delve",
    "testament",
}
ZH_PHRASES = {
    "值得注意的是",
    "综上所述",
    "赋能",
    "降本增效",
    "闭环",
    "无缝",
    "由此可见",
}


def _load(split: str):
    pairs = []
    with open(DATA / f"{split}.jsonl", encoding="utf-8") as f:
        for line in f:
            pairs.append(json.loads(line))
    return pairs


def test_splits_exist_and_nonempty():
    for split in ("train", "val", "test"):
        pairs = _load(split)
        assert len(pairs) > 0, f"{split} is empty"
    assert len(_load("train")) > 10_000


def test_pairs_have_expected_fields():
    for split in ("train", "val", "test"):
        for p in _load(split):
            assert set(p) == {"input_text", "output_text", "lang"}
            assert p["lang"] in {"en", "zh"}
            assert p["input_text"] and p["output_text"]


def test_both_languages_present():
    langs = {p["lang"] for p in _load("train")}
    assert langs == {"en", "zh"}


def test_ai_phrases_removed_in_targets():
    counts = 0
    checked = 0
    for p in _load("train"):
        if "值得注意的是" in p["input_text"] or "it is important to note that" in p["input_text"].lower():
            checked += 1
            target_lower = p["output_text"].lower()
            if p["lang"] == "zh":
                assert "值得注意的是" not in target_lower
                assert "综上所述" not in target_lower
            else:
                assert "it is important to note that" not in target_lower
                assert "it is worth noting that" not in target_lower
            counts += 1
    assert counts > 500


def test_protected_spans_preserved():
    seen = 0
    for p in _load("train") + _load("test"):
        src = p["input_text"]
        if "https://" in src or "ABC-12345" in src or "v2.1.3" in src:
            for token in ("https://api.lynote.ai", "ABC-12345", "v2.1.3"):
                if token in src:
                    assert token in p["output_text"], f"lost {token} in {p['output_text']}"
                    seen += 1
    assert seen > 10


def test_identity_pairs_unchanged():
    for p in _load("train"):
        if p["input_text"] == p["output_text"]:
            assert p["lang"] in {"en", "zh"}
            return
    raise AssertionError("no identity pairs found")


def test_deterministic_generation():
    """Same seed yields the same corpus (checked on a fixed sample)."""
    src = "In today's rapidly evolving world, it is important to note that this solution is robust. Moreover, it is seamless."
    out1 = humanize_rules(src, "en")
    out2 = humanize_rules(src, "en")
    assert out1 == out2


def test_protect_restore_roundtrip():
    text = "Run `pip install x==1.0` on https://example.com/a/b now — the price is $12.50 and \"quoted\" text."
    masked, protected = protect(text)
    assert masked != text
    assert "PROTECTED_0" in masked
    assert restore(masked, protected) == text


def test_rules_en():
    src = "In today's rapidly evolving world, it is important to note that this robust solution serves as a testament to our commitment. Moreover, we leverage cutting-edge technology for a seamless experience."
    out = humanize_rules(src, "en")
    lower = out.lower()
    assert "it is important to note that" not in lower
    assert "moreover" not in lower
    assert "leverage" not in lower
    assert "seamless" not in lower
    assert "robust" not in lower
    assert "testament" not in lower
    assert out != src


def test_rules_zh():
    src = "值得注意的是，在当今快速发展的时代，我们通过赋能团队来实现降本增效，形成完整闭环，无缝连接各个环节。"
    out = humanize_rules(src, "zh")
    assert "值得注意的是" not in out
    assert "赋能" not in out
    assert "降本增效" not in out
    assert "无缝" not in out
    assert out != src


def test_rules_keep_human_prose_intact():
    clean = "I walked the dog this morning and it started raining halfway through the park."
    out = humanize_rules(clean, "en")
    assert out == clean


def test_dataset_reproducibility_from_seed():
    random.seed(7)
    sample1 = random.sample(_load("test"), 50)
    random.seed(7)
    sample2 = random.sample(_load("test"), 50)
    assert [p["input_text"] for p in sample1] == [p["input_text"] for p in sample2]