tests/test_data_ingest.py

from __future__ import annotations

import gzip
import json
import sys
import tempfile
import unittest
from pathlib import Path

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

from turn_detection.data import (  # noqa: E402
    DatasetReadError,
    iter_local_records,
    iter_records,
    read_manifest,
    write_manifest,
)


class LocalIngestionTests(unittest.TestCase):
    def test_jsonl_provenance_and_global_limit(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            first = root / "a.jsonl"
            second = root / "b.jsonl"
            first.write_text('{"id":"a"}\n{"id":"b"}\n', encoding="utf-8")
            second.write_text('{"id":"c"}\n', encoding="utf-8")
            rows = list(iter_records(root, limit=2))
            self.assertEqual([row["id"] for row in rows], ["a", "b"])
            self.assertEqual(rows[1]["__source_row"], 1)
            self.assertTrue(rows[0]["__source_file"].endswith("a.jsonl"))

    def test_gzipped_jsonl(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            path = Path(temporary) / "records.jsonl.gz"
            with gzip.open(path, "wt", encoding="utf-8") as handle:
                handle.write(json.dumps({"id": "compressed"}) + "\n")
            rows = list(iter_local_records(path))
            self.assertEqual(rows[0]["id"], "compressed")

    def test_invalid_json_has_file_and_line_context(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            path = Path(temporary) / "bad.jsonl"
            path.write_text("not json\n", encoding="utf-8")
            with self.assertRaisesRegex(DatasetReadError, "bad.jsonl:1"):
                list(iter_local_records(path))

    def test_manifest_round_trip_is_jsonl_and_atomic(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            path = Path(temporary) / "nested" / "manifest.jsonl"
            count = write_manifest(path, [{"record_id": "एक", "endpoint": True}])
            self.assertEqual(count, 1)
            self.assertEqual(list(read_manifest(path))[0]["record_id"], "एक")
            self.assertFalse(any(path.parent.glob("*.tmp")))


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