doczyai-pipelines/src/tests/test_string_utils.py

import unittest
import json
import pandas as pd
import numpy as np
from src.utils import string_utils
from src.constants.delimiters import Delimiter


class TestStringUtils(unittest.TestCase):
    def test_extract_text_from_delimiters_pipe(self):
        s = "foo |bar| baz |qux|"
        self.assertEqual(
            string_utils.extract_text_from_delimiters(s, Delimiter.PIPE), "qux"
        )
        self.assertEqual(
            string_utils.extract_text_from_delimiters(s, Delimiter.PIPE, 0), "bar"
        )
        self.assertEqual(
            string_utils.extract_text_from_delimiters("no delimiters", Delimiter.PIPE),
            "N/A",
        )
        with self.assertRaises(TypeError):
            string_utils.extract_text_from_delimiters(123, Delimiter.PIPE)
        with self.assertRaises(TypeError):
            string_utils.extract_text_from_delimiters(s, "notadelimiter")

    def test_extract_text_from_delimiters_backtick(self):
        s = "foo `bar` baz `qux`"
        self.assertEqual(
            string_utils.extract_text_from_delimiters(s, Delimiter.BACKTICK), "qux"
        )

    def test_extract_text_from_delimiters_triple_backtick(self):
        s = "foo ```bar``` baz ```qux```"
        self.assertEqual(
            string_utils.extract_text_from_delimiters(s, Delimiter.TRIPLE_BACKTICK),
            "qux",
        )

    def test_page_key_sort(self):
        keys = ["10", "2", "A", "1", "B", "1.5"]
        sorted_keys = sorted(keys, key=string_utils.page_key_sort)
        self.assertEqual(sorted_keys, ["1", "1.5", "2", "10", "A", "B"])

    def test_json_parsing_search(self):
        s = '{"FIELD1": "value1", "FIELD2": "value2"}'
        fields = ["FIELD1", "FIELD2"]
        result = string_utils.json_parsing_search(s, fields)
        self.assertEqual(result, {"FIELD1": "value1", "FIELD2": "value2"})

    def test_universal_json_load_valid(self):
        # Simple dict
        s = '{"a": 1, "b": 2}'
        self.assertEqual(string_utils.universal_json_load(s), {"a": 1, "b": 2})
        # Simple list
        s = "[1, 2, 3]"
        self.assertEqual(string_utils.universal_json_load(s), [1, 2, 3])
        # Nested dict
        s = '{"a": {"b": [1, 2, {"c": "d"}]}}'
        self.assertEqual(
            string_utils.universal_json_load(s), {"a": {"b": [1, 2, {"c": "d"}]}}
        )
        # List of dicts
        s = '[{"x": 1}, {"y": 2}]'
        self.assertEqual(string_utils.universal_json_load(s), [{"x": 1}, {"y": 2}])
        # Dict of lists
        s = '{"a": [1, 2], "b": [3, 4]}'
        self.assertEqual(
            string_utils.universal_json_load(s), {"a": [1, 2], "b": [3, 4]}
        )
        # List of strings
        s = '["foo", "bar", "baz"]'
        self.assertEqual(string_utils.universal_json_load(s), ["foo", "bar", "baz"])
        # JSON with text before/after
        s = 'start {"field": "value"} end'
        self.assertEqual(string_utils.universal_json_load(s), {"field": "value"})
        s = "before [1, 2, 3] after"
        self.assertEqual(string_utils.universal_json_load(s), [1, 2, 3])
        # Multiple objects, last wins
        s = 'first {"a": 1} middle {"b": 2} last {"c": 3}'
        self.assertEqual(string_utils.universal_json_load(s), {"c": 3})
        # Multiple arrays, last wins
        s = "first [1] middle [2] last [3]"
        self.assertEqual(string_utils.universal_json_load(s), [3])
        # Nested JSON with text 1
        s = 'junk {"outer": {"inner": [1, 2, {"deep": true}]}} more junk'
        self.assertEqual(
            string_utils.universal_json_load(s),
            {"outer": {"inner": [1, 2, {"deep": True}]}},
        )
        # Nested JSON with text 2
        s = 'junk [{"dict1_key": "dict1_val"}, {"dict2_key" : "dict2_val"}] more junk'
        self.assertEqual(
            string_utils.universal_json_load(s),
            [{"dict1_key": "dict1_val"}, {"dict2_key": "dict2_val"}],
        )
        # List of dicts
        s = 'junk```json[{"a1": "value_a1","b1": "value_b1"},{"a2": "value_a2","b2": "value_b2"}]```'
        self.assertEqual(
            string_utils.universal_json_load(s),
            [
                {"a1": "value_a1", "b1": "value_b1"},
                {"a2": "value_a2", "b2": "value_b2"},
            ],
        )
        # Boolean and null values
        s = '{"ok": true, "val": null}'
        self.assertEqual(string_utils.universal_json_load(s), {"ok": True, "val": None})

    def test_universal_json_load_invalid(self):
        # No valid JSON
        s = "no json here"
        with self.assertRaises(ValueError):
            string_utils.universal_json_load(s)
        # Broken JSON
        s = '{broken: "json"}'
        with self.assertRaises(ValueError):
            string_utils.universal_json_load(s)
        # Non-string input
        with self.assertRaises(TypeError):
            string_utils.universal_json_load(123)
        # Unbalanced brackets
        s = '{"a": [1, 2, 3}'
        with self.assertRaises(ValueError):
            string_utils.universal_json_load(s)
        # Empty string
        s = ""
        with self.assertRaises(ValueError):
            string_utils.universal_json_load(s)

    def test_universal_json_load_edge_cases(self):
        # Empty object
        s = "text {} more text"
        self.assertEqual(string_utils.universal_json_load(s), {})
        # Empty array
        s = "text [] more text"
        self.assertEqual(string_utils.universal_json_load(s), [])
        # Large object
        large = {f"k{i}": i for i in range(100)}
        s = f"before {json.dumps(large)} after"
        self.assertEqual(string_utils.universal_json_load(s), large)

    def test_contains_reimbursement(self):
        s = "90% reimbursement of billed charges"
        self.assertTrue(string_utils.contains_reimbursement(s, method="keyword"))
        s = "Payment of $500 per visit"
        self.assertTrue(string_utils.contains_reimbursement(s, method="regex"))
        s = "No reimbursement here"
        self.assertTrue(string_utils.contains_reimbursement(s, method="keyword"))
        with self.assertRaises(ValueError):
            string_utils.contains_reimbursement(s, method="invalid")

    def test_is_empty_string(self):
        self.assertTrue(string_utils.is_empty(""))
        self.assertTrue(string_utils.is_empty("N/A"))
        self.assertFalse(string_utils.is_empty("valid"))

    def test_is_empty_list(self):
        self.assertTrue(string_utils.is_empty([]))
        self.assertTrue(string_utils.is_empty([None, "", "N/A"]))
        self.assertFalse(string_utils.is_empty(["valid"]))

    def test_is_empty_pandas_series(self):
        s = pd.Series(["valid", "", None, "N/A"])
        mask = string_utils.is_empty(s, pd_mask=True)
        self.assertTrue(mask.iloc[1])
        self.assertTrue(mask.iloc[2])
        self.assertTrue(mask.iloc[3])
        self.assertFalse(mask.iloc[0])
        self.assertFalse(
            string_utils.is_empty(pd.Series(["valid", None]), pd_mask=False)
        )
        self.assertTrue(
            string_utils.is_empty(pd.Series([None, "", "N/A"]), pd_mask=False)
        )

    def test_datetime_str(self):
        result = string_utils.datetime_str()
        self.assertTrue(result.startswith("["))
        self.assertTrue(result.endswith("]"))

    def test_extract_signature_page(self):
        text_dict = {
            "1": "Contract page 1",
            "2": "This page has 2 signatures at the bottom",
            "3": "Final page",
        }
        result = string_utils.extract_signature_page(text_dict, "test.pdf")
        self.assertIn("2", result)
        self.assertEqual(result["2"], "This page has 2 signatures at the bottom")
        text_dict = {
            "1": "Contract page 1",
            "2": "In Witness Whereof",
            "3": "Final page",
        }
        result = string_utils.extract_signature_page(text_dict, "test.pdf")
        self.assertIn("2", result)
        self.assertEqual(result["2"], "In Witness Whereof")
        text_dict = {
            "1": "Contract page 1",
            "2": "Signature page follows this section",
            "3": "signature page",
        }
        result = string_utils.extract_signature_page(text_dict, "test.pdf")
        self.assertIn("3", result)

    def test_extract_effective_date_pages(self):
        text_dict = {
            "1": "Contract effective date: January 1, 2024",
            "2": "General terms and conditions",
            "3": "This page has 1 signature",
        }
        result = string_utils.extract_effective_date_pages(text_dict, "test.pdf")
        self.assertIn("1", result)
        self.assertIn("3", result)
        self.assertEqual(len(result), 2)

    def test_normalize_state_to_abbreviation(self):
        self.assertEqual(
            string_utils.normalize_state_to_abbreviation("California"), "CA"
        )
        self.assertEqual(string_utils.normalize_state_to_abbreviation("ca"), "CA")
        self.assertEqual(string_utils.normalize_state_to_abbreviation("N.Y."), "NY")
        self.assertEqual(
            string_utils.normalize_state_to_abbreviation("Unknown State"),
            "Unknown State",
        )
        self.assertEqual(string_utils.normalize_state_to_abbreviation("N/A"), "N/A")

    def test_flatten_to_strings(self):
        self.assertEqual(
            string_utils.flatten_to_strings(["a", ["b", "c"], "d"]),
            ["a", "b", "c", "d"],
        )
        self.assertEqual(string_utils.flatten_to_strings("single"), ["single"])
        self.assertEqual(
            string_utils.flatten_to_strings([1, [2, 3], "text"]),
            ["1", "2", "3", "text"],
        )
        self.assertEqual(
            string_utils.flatten_to_strings([" text ", ["  nested  "]]),
            ["text", "nested"],
        )

    def test_universal_json_load_multiple_invalid(self):
        # Multiple invalid JSON objects
        s = "{broken: true} some text {not: valid} more text"
        with self.assertRaises(ValueError):
            string_utils.universal_json_load(s)

    def test_universal_json_load_one_invalid_multiple_valid(self):
        # One invalid, multiple valid (should take the last valid)
        s = '{broken: true} before {"a": 1} middle {"b": 2} end'
        self.assertEqual(string_utils.universal_json_load(s), {"b": 2})

    def test_universal_json_load_first_nested(self):
        # Multiple, with the first one being nested
        s = '{"outer": {"inner": [1, 2]}} middle {"simple": true} end'
        self.assertEqual(string_utils.universal_json_load(s), {"simple": True})

    def test_universal_json_load_last_nested(self):
        # Multiple, with the last one being nested
        s = '{"simple": true} middle {"outer": {"inner": [1, 2]}} end'
        self.assertEqual(
            string_utils.universal_json_load(s), {"outer": {"inner": [1, 2]}}
        )

    def test_universal_json_load_brackets_in_strings(self):
        # Brackets inside string values should be ignored during bracket matching
        s = '{"text": "array [1] and {2}"}'
        self.assertEqual(
            string_utils.universal_json_load(s), {"text": "array [1] and {2}"}
        )
        # Double brackets in strings (the original bug case)
        s = '{"field": "Exhibit [[ to Attachment ]"}'
        self.assertEqual(
            string_utils.universal_json_load(s), {"field": "Exhibit [[ to Attachment ]"}
        )
        # Unbalanced brackets in strings
        s = '{"text": "missing close [ bracket"}'
        self.assertEqual(
            string_utils.universal_json_load(s), {"text": "missing close [ bracket"}
        )
        # Curly braces in strings
        s = '{"text": "object {key: value} here"}'
        self.assertEqual(
            string_utils.universal_json_load(s), {"text": "object {key: value} here"}
        )
        # Escaped quotes should not break string detection
        s = '{"text": "say \\"hello\\""}'
        self.assertEqual(string_utils.universal_json_load(s), {"text": 'say "hello"'})
        # Escaped quotes with brackets
        s = '{"text": "say \\"hello [world]\\"" }'
        self.assertEqual(
            string_utils.universal_json_load(s), {"text": 'say "hello [world]"'}
        )
        # Real-world LLM response case with markdown and brackets in JSON strings
        raw_input = 'Here is the analysis:\n\n```json\n[{"SERVICE_TERM": "Inpatient services", "REIMB_TERM": "100% of rates in Exhibit [[ to Attachment ]"}, {"SERVICE_TERM": "Outpatient services", "REIMB_TERM": "63% of rates in Exhibit [[ to Attachment ]"}]\n```'
        result = string_utils.universal_json_load(raw_input)
        self.assertEqual(len(result), 2)
        self.assertEqual(result[0]["SERVICE_TERM"], "Inpatient services")
        self.assertIn("Exhibit [[", result[0]["REIMB_TERM"])


if __name__ == "__main__":
    unittest.main()