"""Tests for tolerant JSON object extraction from model output.""" from __future__ import annotations import unittest from src.utils.json_repair import parse_json_object class JsonRepairTest(unittest.TestCase): def test_parses_complete_object_with_surrounding_text(self) -> None: payload = parse_json_object('Here is the archive:\n{"name": "mug"}\nDone.') self.assertEqual(payload, {"name": "mug"}) def test_repairs_missing_outer_closing_brace(self) -> None: payload = parse_json_object( """ { "persona": {"object_name": "coffee mug"}, "diary": {"title": "Secret Diary - Day 310"} """ ) self.assertEqual(payload["persona"]["object_name"], "coffee mug") self.assertEqual(payload["diary"]["title"], "Secret Diary - Day 310") def test_does_not_repair_unterminated_string(self) -> None: with self.assertRaises(ValueError): parse_json_object('{"name": "coffee mug}') if __name__ == "__main__": unittest.main()