import argparse import json import shutil import subprocess import tempfile import unittest from pathlib import Path from scripts import generate_training_data as bridge class GenerateTrainingDataTests(unittest.TestCase): def setUp(self): self.repo = Path(__file__).resolve().parents[1] self.temp = tempfile.TemporaryDirectory() self.root = Path(self.temp.name) self.dictionary_dir = self.root / "dictionary" self.dictionary_dir.mkdir() shutil.copy2(self.repo / "dictionary" / "static-base.json", self.dictionary_dir / "static-base.json") subprocess.run( [ "cargo", "run", "--", "dictionary", "generate", "--out-dir", str(self.dictionary_dir), ], cwd=self.repo, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) self.output_dir = self.root / "workspace_Data" / "data" self.output_dir.mkdir(parents=True) self.manifest = self.root / "workspace_Data" / "manifest.llm.json" self.manifest.write_text(json.dumps({"manifest_version": "0.0.0", "datasets": {}})) def tearDown(self): self.temp.cleanup() def args(self, **overrides): values = { "dictionary_dir": str(self.dictionary_dir), "output_dir": str(self.output_dir), "manifest": str(self.manifest), "dry_run": False, "sdg_host": None, "sdg_user": None, "sdg_key": None, "sdg_model_alias": "nvidia-text", "num_records": 1000, "timestamp": 1234567890, } values.update(overrides) return argparse.Namespace(**values) def test_seed_generation_covers_current_dictionary(self): artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir) records = bridge.build_seed_records(artifacts, generated_at=123) actions = {record["metadata"]["action_id"] for record in records} self.assertEqual(len(actions), 13) self.assertIn("calculate", actions) self.assertIn("delete_file", actions) self.assertEqual(len(records), 260) def test_compact_output_uses_codebook_codes(self): artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir) compact = bridge.compact_target( artifacts, "hash_string", {"text": "hello", "algorithm": "sha256"}, ) self.assertEqual(compact["a"], artifacts["action_to_code"]["hash_string"]) self.assertIn(artifacts["field_to_code"]["text"], compact["p"]) self.assertIn(artifacts["field_to_code"]["algorithm"], compact["p"]) def test_empty_payload_action(self): artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir) records = bridge.build_seed_records(artifacts, generated_at=123) uuid_records = [ record for record in records if record["metadata"]["action_id"] == "generate_uuid" ] self.assertTrue(uuid_records) output = json.loads(uuid_records[0]["output"]) compact = json.loads(uuid_records[0]["metadata"]["compact"]) self.assertEqual(output["payload"], {}) self.assertEqual(compact["p"], {}) def test_split_includes_every_action(self): artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir) records = bridge.build_seed_records(artifacts, generated_at=123) train, validation = bridge.split_records_by_action(records) train_actions = {record["metadata"]["action_id"] for record in train} validation_actions = {record["metadata"]["action_id"] for record in validation} self.assertEqual(train_actions, validation_actions) self.assertEqual(len(train_actions), 13) def test_dry_run_writes_nothing(self): summary = bridge.run_bridge(self.args(dry_run=True)) self.assertTrue(summary["dry_run"]) self.assertFalse((self.output_dir / "module_controller_intents_1234567890").exists()) manifest = json.loads(self.manifest.read_text()) self.assertEqual(manifest["datasets"], {}) def test_manifest_patch_adds_dataset_without_damaging_existing_keys(self): bridge.run_bridge(self.args()) manifest = json.loads(self.manifest.read_text()) self.assertEqual(manifest["manifest_version"], "0.0.0") dataset = manifest["datasets"]["module_controller_intents_1234567890"] self.assertEqual(dataset["schema_profile"], "instruction") self.assertEqual(dataset["codebook_version"], "1.0.0") self.assertIn("codebook_checksum", dataset) self.assertTrue((self.output_dir / "module_controller_intents_1234567890" / "train.jsonl").exists()) def test_training_data_relative_path_prefers_custom_root_shape(self): path = Path("/Volumes/Zoe/custom-local-llm/training-data/module_controller_intents_123") self.assertEqual( bridge.relative_data_path(path), "training-data/module_controller_intents_123", ) def test_data_designer_config_declares_structured_expansion(self): config = bridge.data_designer_config( dataset_key="module_controller_intents_123", generated_at=123, remote_seed="/tmp/seed.jsonl", action_ids=["logger"], model_alias="nvidia-text", ) self.assertIn("load_config_builder", config) self.assertIn('column_type="llm-structured"', config) self.assertIn('model_alias="nvidia-text"', config) self.assertIn("expanded_record", config) def test_remote_data_designer_command_sources_mimir_env(self): command = bridge.remote_data_designer_create_command( remote_config="/tmp/controller.py", num_records=42, dataset_key="module_controller_intents_123", artifact_path="/tmp/artifacts", ) self.assertIn(bridge.MIMIR_DATA_DESIGNER_ENV, command) self.assertIn(bridge.MIMIR_DATA_DESIGNER_BIN, command) self.assertIn("--num-records 42", command) self.assertIn("--no-tui", command) def test_remote_bash_command_quotes_inner_command(self): command = bridge.remote_bash_command("set -a; echo ready") self.assertTrue(command.startswith("bash -lc ")) self.assertIn("'set -a; echo ready'", command) def test_normalize_expanded_jsonl_extracts_expanded_record(self): raw = self.root / "expanded.raw.jsonl" normalized = self.root / "expanded.jsonl" source_row = { "expanded_record": { "instruction": bridge.INSTRUCTION, "input": "record this message", "output": json.dumps({"action_id": "logger", "payload": {"message": "hello"}}), "metadata": { "action_id": "logger", "source": "data_designer", "split": "train", "generated_at": 123, }, } } bridge.write_jsonl(raw, [source_row]) bridge.normalize_expanded_jsonl(raw, normalized) rows = bridge.read_jsonl(normalized) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["metadata"]["action_id"], "logger") if __name__ == "__main__": unittest.main()