| """Tests for non-UI dataset and trace export tooling.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import tempfile |
| import unittest |
| from collections import Counter |
| from pathlib import Path |
|
|
| from scripts.export_traces import export_trace_jsonl |
| from scripts.generate_dataset import build_sft_records, write_sft_jsonl |
| from scripts.generate_sample_traces import generate_sample_traces |
| from scripts.prepare_curated_dataset import MODES, build_curated_records, write_jsonl |
| from src.models.schema import TraceRecord |
|
|
|
|
| class DatasetToolingTest(unittest.TestCase): |
| def test_build_sft_records(self) -> None: |
| records = build_sft_records(7) |
| assistant_payload = json.loads(records[0]["messages"][2]["content"]) |
|
|
| self.assertEqual(len(records), 7) |
| self.assertEqual(records[0]["split"], "preview") |
| self.assertIn("object_understanding", records[0]) |
| self.assertIn("persona", assistant_payload) |
| self.assertIn("diary", assistant_payload) |
|
|
| def test_write_sft_jsonl(self) -> None: |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| output_path = Path(tmp_dir) / "preview.jsonl" |
| write_sft_jsonl(build_sft_records(3), output_path) |
| rows = [ |
| json.loads(line) |
| for line in output_path.read_text(encoding="utf-8").splitlines() |
| ] |
|
|
| self.assertEqual(len(rows), 3) |
| self.assertEqual(rows[0]["id"], "sft-preview-0001") |
|
|
| def test_build_curated_records(self) -> None: |
| records = build_curated_records(10) |
| assistant_payload = json.loads(records[0]["messages"][2]["content"]) |
|
|
| self.assertEqual(len(records), 10) |
| self.assertEqual(records[0]["split"], "train") |
| self.assertEqual(records[0]["source"], "objectverse-diary-synthetic-curated-v1") |
| self.assertIn("curation_notes", records[0]) |
| self.assertIn("persona", assistant_payload) |
| self.assertIn("diary", assistant_payload) |
|
|
| def test_build_curated_v2_records_has_broader_balanced_coverage(self) -> None: |
| records = build_curated_records(200, version="v2") |
| object_names = [ |
| record["object_understanding"]["object"]["name"] |
| for record in records |
| ] |
| mode_counts = Counter(record["mode"] for record in records) |
| object_mode_pairs = {(name, record["mode"]) for name, record in zip(object_names, records)} |
| assistant_payload = json.loads(records[0]["messages"][2]["content"]) |
|
|
| self.assertEqual(len(records), 200) |
| self.assertEqual(records[0]["source"], "objectverse-diary-synthetic-curated-v2") |
| self.assertGreaterEqual(len(set(object_names)), 40) |
| self.assertEqual(mode_counts, Counter({mode: 40 for mode in MODES})) |
| self.assertEqual(len(object_mode_pairs), 200) |
| self.assertIn("scene_detail", records[0]) |
| self.assertIn("persona", assistant_payload) |
| self.assertIn("diary", assistant_payload) |
|
|
| def test_write_curated_jsonl(self) -> None: |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| output_path = Path(tmp_dir) / "curated.jsonl" |
| write_jsonl(build_curated_records(2), output_path) |
| rows = [ |
| json.loads(line) |
| for line in output_path.read_text(encoding="utf-8").splitlines() |
| ] |
|
|
| self.assertEqual(len(rows), 2) |
| self.assertEqual(rows[0]["id"], "curated-synthetic-0001") |
|
|
| def test_export_trace_jsonl(self) -> None: |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| sample_dir = Path(tmp_dir) / "samples" |
| output_path = Path(tmp_dir) / "public_traces.jsonl" |
| generate_sample_traces(sample_dir) |
| count = export_trace_jsonl(sample_dir, output_path) |
| rows = [ |
| json.loads(line) |
| for line in output_path.read_text(encoding="utf-8").splitlines() |
| ] |
|
|
| self.assertEqual(count, 6) |
| self.assertEqual(len(rows), 6) |
| TraceRecord.model_validate(rows[0]) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
