File size: 2,287 Bytes
35d483e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | 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()
|