tests/test_json_repair.py

"""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()