doczyai-pipelines/fieldExtraction/tests/test_string_utils.py

import unittest
import json
import pandas as pd
import numpy as np
from src.utils import string_utils
from 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": "Signature page follows this section", "3": "Final signature authorization page"}
        result = string_utils.extract_signature_page(text_dict, "test.pdf")
        self.assertIn("2", result)
        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]}})

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