Spaces:
Running
Running
| import json | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| from llm.training_data import load_emotion_ids, load_examples, sha256_file | |
| class TrainingDataTests(unittest.TestCase): | |
| def test_loads_valid_examples_and_honors_row_limit(self) -> None: | |
| with tempfile.TemporaryDirectory() as directory: | |
| path = Path(directory) / "data.jsonl" | |
| path.write_text( | |
| "\n".join( | |
| json.dumps(row) | |
| for row in ( | |
| {"user": "Hello", "emotion": "happy", "reply": "Hi!"}, | |
| {"user": "Wait", "emotion": "calm", "reply": "I can wait."}, | |
| ) | |
| ) | |
| + "\n" | |
| ) | |
| rows = load_examples(path, {"happy", "calm"}, max_rows=1) | |
| self.assertEqual(rows, [{"user": "Hello", "emotion": "happy", "reply": "Hi!"}]) | |
| self.assertEqual(len(sha256_file(path)), 64) | |
| def test_rejects_unknown_emotion(self) -> None: | |
| with tempfile.TemporaryDirectory() as directory: | |
| path = Path(directory) / "data.jsonl" | |
| path.write_text('{"user":"Hello","emotion":"unsafe","reply":"Hi"}\n') | |
| with self.assertRaisesRegex(ValueError, "unknown emotion"): | |
| load_examples(path, {"happy"}) | |
| def test_loads_unique_emotion_ids(self) -> None: | |
| with tempfile.TemporaryDirectory() as directory: | |
| path = Path(directory) / "emotions.json" | |
| path.write_text('[{"id":"happy"},{"id":"calm"}]') | |
| self.assertEqual(load_emotion_ids(path), ["happy", "calm"]) | |
| path.write_text('[{"id":"happy"},{"id":"happy"}]') | |
| with self.assertRaisesRegex(ValueError, "must be unique"): | |
| load_emotion_ids(path) | |
| if __name__ == "__main__": | |
| unittest.main() | |