tests/test_data_prep.py

"""Verify reform output: messages serialize as clean structs (not arrow.json
strings), and every row is a well-formed [system, user, assistant] SFT example."""

import json
import pathlib

import pyarrow.parquet as pq
import tqdm
from datasets import load_dataset
from datasets.utils.logging import disable_progress_bar

# The HA pytest plugin fails teardown on any lingering thread; datasets' tqdm bar
# spawns a monitor thread. Suppress both so the check stays clean.
disable_progress_bar()
tqdm.tqdm.monitor_interval = 0

_DATA_DIR = pathlib.Path(__file__).resolve().parent.parent / "data" / "processed" / "ha_actions"
_SPLITS = ("train", "test")


def _check_split(split: str) -> None:
    path = _DATA_DIR / f"{split}.parquet"

    # 1. on-disk schema is a struct list, NOT the arrow.json extension.
    field = pq.read_schema(path).field("messages")
    assert "struct" in str(field.type), f"{split}: expected struct list, got {field.type}"
    assert "json" not in str(field.type).lower(), f"{split}: arrow.json extension leaked back in"

    # 2. loaded rows are native dicts, well-formed, with a parseable assistant payload.
    ds = load_dataset("parquet", data_files=str(path), split="train")
    for r in ds:
        msgs = r["messages"]
        assert [m["role"] for m in msgs] == ["system", "user", "assistant"], msgs
        assert all(isinstance(m["content"], str) for m in msgs)
        payload = json.loads(msgs[2]["content"])  # must be valid JSON
        assert "intents" in payload or "response" in payload, payload
    print(f"[ok] {split}: {len(ds)} rows, struct schema, all assistant turns valid JSON")


def test_reform_output() -> None:
    for split in _SPLITS:
        _check_split(split)


if __name__ == "__main__":
    test_reform_output()