File size: 1,064 Bytes
dd6cefc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
"""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()