tests/test_data_cli.py

from __future__ import annotations

import json
import struct
import subprocess
import sys
import tempfile
import unittest
import wave
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from turn_detection.data import read_manifest  # noqa: E402


def write_wav(path: Path, value: int) -> None:
    with wave.open(str(path), "wb") as handle:
        handle.setnchannels(1)
        handle.setsampwidth(2)
        handle.setframerate(16_000)
        handle.writeframes(struct.pack("<h", value) * 160)


class DataCliTests(unittest.TestCase):
    def test_audit_then_prepare_splits_without_optional_dependencies(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            for index in range(6):
                write_wav(root / f"{index}.wav", index + 1)
            source = root / "records.jsonl"
            with source.open("w", encoding="utf-8") as handle:
                for index in range(6):
                    handle.write(
                        json.dumps(
                            {
                                "audio": {"path": f"{index}.wav"},
                                "id": f"id-{index}",
                                "language": "hi" if index % 2 else "en",
                                "endpoint_bool": bool(index % 2),
                                "midfiller": None,
                                "endfiller": False,
                                "synthetic": False,
                                "dataset": "test",
                            }
                        )
                        + "\n"
                    )
            manifest = root / "manifest.jsonl"
            audit_report = root / "audit.json"
            audit = subprocess.run(
                [
                    sys.executable,
                    str(PROJECT_ROOT / "scripts" / "audit_dataset.py"),
                    str(source),
                    "--output",
                    str(manifest),
                    "--report",
                    str(audit_report),
                    "--progress-every",
                    "0",
                ],
                cwd=PROJECT_ROOT,
                capture_output=True,
                text=True,
                check=False,
            )
            self.assertEqual(audit.returncode, 0, audit.stderr)
            self.assertEqual(len(list(read_manifest(manifest))), 6)

            splits = root / "splits.jsonl"
            split_report = root / "splits.json"
            prepare = subprocess.run(
                [
                    sys.executable,
                    str(PROJECT_ROOT / "scripts" / "prepare_splits.py"),
                    "--input",
                    str(manifest),
                    "--output",
                    str(splits),
                    "--report",
                    str(split_report),
                    "--split",
                    "train=2",
                    "--split",
                    "validation=1",
                    "--limit",
                    "5",
                ],
                cwd=PROJECT_ROOT,
                capture_output=True,
                text=True,
                check=False,
            )
            self.assertEqual(prepare.returncode, 0, prepare.stderr)
            split_rows = list(read_manifest(splits))
            self.assertEqual(len(split_rows), 5)
            self.assertTrue(all(row.get("split") for row in split_rows))


if __name__ == "__main__":
    unittest.main()