Merged in feature/black-and-isort (pull request #717)

Feature/black and isort

* black and isort constants

* black and isort codes

* black and isort crosswalks

* black and isort investment

* black and isort prompts

* isort testbed

* isort tracking

* black and isort utils

* poetry and black tests


Approved-by: Katon Minhas
This commit is contained in:
Alex Galarce
2025-09-30 20:56:02 +00:00
parent ae4eda6f25
commit 674add9a3e
43 changed files with 841 additions and 515 deletions
-1
View File
@@ -4,7 +4,6 @@ import unittest
from unittest.mock import MagicMock, patch
import pandas as pd
import src.codes.code_funcs as code_funcs
from constants.constants import Constants
from constants.delimiters import Delimiter
@@ -2,7 +2,6 @@ import json
import pandas as pd
import pytest
from src.crosswalk.crosswalk_builder import CrosswalkBuilder
from src.utils.crosswalk_utils import apply_crosswalk
+221 -114
View File
@@ -1,13 +1,12 @@
import unittest
import json
import os
from unittest.mock import patch, MagicMock
import unittest
from unittest.mock import MagicMock, patch
import pandas as pd
from constants.constants import Constants
from src.investment import dynamic_funcs, prompt_calls
from src.prompts.fieldset import Field, FieldSet
from constants.constants import Constants
class TestDynamicFuncsWithRealConstants(unittest.TestCase):
@@ -17,81 +16,107 @@ class TestDynamicFuncsWithRealConstants(unittest.TestCase):
"""Set up test fixtures before each test method."""
# Create real Constants object
self.constants = Constants()
# Create test fields
self.test_field = Field.from_values(
field_name="TEST_FIELD",
relationship="one-to-one",
field_type="exhibit_level",
prompt="What is the test field?",
base_field="TEST_FIELD_BASE"
base_field="TEST_FIELD_BASE",
)
self.test_field2 = Field.from_values(
field_name="TEST_FIELD2",
relationship="one-to-many",
field_type="exhibit_level",
prompt="What is the second test field?",
base_field="TEST_FIELD2_BASE"
base_field="TEST_FIELD2_BASE",
)
self.escalator_field = Field.from_values(
field_name="RATE_ESCALATOR_DESC",
relationship="one-to-many",
field_type="rate_escalator",
prompt="What is the rate escalator?",
base_field="RATE_ESCALATOR_STATEMENT"
base_field="RATE_ESCALATOR_STATEMENT",
)
self.outlier_field = Field.from_values(
field_name="OUTLIER_TERMS",
relationship="one-to-many",
field_type="exhibit_level",
prompt="What are the outlier terms?",
base_field="OUTLIER_TERMS"
base_field="OUTLIER_TERMS",
)
# Create test FieldSets
self.test_fieldset = FieldSet()
self.test_fieldset.add_field(self.test_field)
self.test_fieldset.add_field(self.test_field2)
# Mock filename
self.filename = "test_file.pdf"
# Sample text
self.exhibit_text = "This is a sample exhibit text with test data."
self.exhibit_header = "This is a sample exhibit header with test data."
# Sample answer dictionaries
self.answer_dicts = [
{"TEST_FIELD": "Value 1", "TEST_FIELD2": "", "RATE_ESCALATOR_DESC": "", "AARETE_DERIVED_LOB": "Commercial"},
{"TEST_FIELD": "Value 2", "TEST_FIELD2": "", "RATE_ESCALATOR_DESC": "", "AARETE_DERIVED_LOB": "Medicare"},
{"TEST_FIELD": "Value 3", "TEST_FIELD2": "", "OUTLIER_TERMS": "", "AARETE_DERIVED_LOB": ""}
{
"TEST_FIELD": "Value 1",
"TEST_FIELD2": "",
"RATE_ESCALATOR_DESC": "",
"AARETE_DERIVED_LOB": "Commercial",
},
{
"TEST_FIELD": "Value 2",
"TEST_FIELD2": "",
"RATE_ESCALATOR_DESC": "",
"AARETE_DERIVED_LOB": "Medicare",
},
{
"TEST_FIELD": "Value 3",
"TEST_FIELD2": "",
"OUTLIER_TERMS": "",
"AARETE_DERIVED_LOB": "",
},
]
@patch('src.utils.llm_utils.invoke_claude')
@patch('src.utils.string_utils.universal_json_load')
@patch("src.utils.llm_utils.invoke_claude")
@patch("src.utils.string_utils.universal_json_load")
def test_prompt_dynamic(self, mock_json_load, mock_invoke_claude):
"""Test the prompt_dynamic function."""
# Setup
field_prompts = {
"TEST_FIELD": "What is the test field?",
"TEST_FIELD2": "What is the second test field?"
"TEST_FIELD2": "What is the second test field?",
}
# Configure mocks
mock_invoke_claude.return_value = '{"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}'
mock_json_load.return_value = {"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}
mock_invoke_claude.return_value = (
'{"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}'
)
mock_json_load.return_value = {
"TEST_FIELD": ["Value 1"],
"TEST_FIELD2": ["Value 2"],
}
# Execute
result = prompt_calls.prompt_dynamic(self.exhibit_text, field_prompts, self.filename)
result = prompt_calls.prompt_dynamic(
self.exhibit_text, field_prompts, self.filename
)
# Assert
mock_invoke_claude.assert_called_once()
mock_json_load.assert_called_once_with('{"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}')
self.assertEqual(result, {"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]})
mock_json_load.assert_called_once_with(
'{"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}'
)
self.assertEqual(
result, {"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}
)
# Verify that the prompt contains our text
prompt_call = mock_invoke_claude.call_args[0][0]
self.assertIn(self.exhibit_text, prompt_call)
@@ -103,10 +128,12 @@ class TestDynamicFuncsWithRealConstants(unittest.TestCase):
# Setup
one_to_one_fields = FieldSet()
field_to_add = self.test_field
# Execute
result = dynamic_funcs.add_full_context_field(one_to_one_fields, field_to_add, self.answer_dicts)
result = dynamic_funcs.add_full_context_field(
one_to_one_fields, field_to_add, self.answer_dicts
)
# Assert
self.assertEqual(len(result.fields), 1)
self.assertEqual(result.fields[0].field_name, "TEST_FIELD")
@@ -122,12 +149,14 @@ class TestDynamicFuncsWithRealConstants(unittest.TestCase):
relationship="one-to-one",
field_type="exhibit_level",
prompt="What is the program?",
base_field="PROGRAM"
base_field="PROGRAM",
)
# Execute - use our test answer_dicts which has multiple non-empty LOBs
result = dynamic_funcs.add_full_context_field(one_to_one_fields, program_field, self.answer_dicts)
result = dynamic_funcs.add_full_context_field(
one_to_one_fields, program_field, self.answer_dicts
)
# Assert
self.assertEqual(len(result.fields), 0) # Field should not be added
@@ -140,76 +169,102 @@ class TestDynamicFuncsWithRealConstants(unittest.TestCase):
relationship="one-to-one",
field_type="exhibit_level",
prompt="What is the program?",
base_field="PROGRAM"
base_field="PROGRAM",
)
# Use modified answer_dicts with only one non-empty LOB
single_lob_answer_dicts = [
{"TEST_FIELD": "Value 1", "TEST_FIELD2": "", "RATE_ESCALATOR_DESC": "", "AARETE_DERIVED_LOB": "Commercial"},
{"TEST_FIELD": "Value 2", "TEST_FIELD2": "", "RATE_ESCALATOR_DESC": "", "AARETE_DERIVED_LOB": ""},
{"TEST_FIELD": "Value 3", "TEST_FIELD2": "", "OUTLIER_TERMS": "", "AARETE_DERIVED_LOB": ""}
{
"TEST_FIELD": "Value 1",
"TEST_FIELD2": "",
"RATE_ESCALATOR_DESC": "",
"AARETE_DERIVED_LOB": "Commercial",
},
{
"TEST_FIELD": "Value 2",
"TEST_FIELD2": "",
"RATE_ESCALATOR_DESC": "",
"AARETE_DERIVED_LOB": "",
},
{
"TEST_FIELD": "Value 3",
"TEST_FIELD2": "",
"OUTLIER_TERMS": "",
"AARETE_DERIVED_LOB": "",
},
]
# Execute
result = dynamic_funcs.add_full_context_field(one_to_one_fields, program_field, single_lob_answer_dicts)
result = dynamic_funcs.add_full_context_field(
one_to_one_fields, program_field, single_lob_answer_dicts
)
# Assert
self.assertEqual(len(result.fields), 1) # Field should be added
self.assertEqual(result.fields[0].field_name, "PROGRAM")
@patch('src.investment.dynamic_funcs.Field.load_from_file')
@patch('src.investment.dynamic_funcs.FieldSet')
def test_get_dynamic_one_to_one_fields_any_empty(self, mock_fieldset_class, mock_load_field):
@patch("src.investment.dynamic_funcs.Field.load_from_file")
@patch("src.investment.dynamic_funcs.FieldSet")
def test_get_dynamic_one_to_one_fields_any_empty(
self, mock_fieldset_class, mock_load_field
):
"""Test getting dynamic one-to-one fields with any empty criterion."""
# Setup
# Configure mock for fieldset class
mock_all_empty_fields = MagicMock()
mock_all_empty_fields.fields = []
mock_any_empty_fields = MagicMock()
mock_field1 = MagicMock()
mock_field1.field_name = "RATE_ESCALATOR_DESC"
mock_field1.base_field = "RATE_ESCALATOR_STATEMENT"
mock_any_empty_fields.fields = [mock_field1]
mock_one_to_one_fields = MagicMock()
mock_one_to_one_fields.fields = []
mock_fieldset_class.side_effect = [
mock_all_empty_fields,
mock_any_empty_fields,
mock_one_to_one_fields
mock_one_to_one_fields,
]
# Configure mock for loading fields
escalator_statement_field = Field.from_values(
field_name="RATE_ESCALATOR_STATEMENT",
relationship="one-to-one",
field_type="full_context",
prompt="What is the rate escalator statement?"
prompt="What is the rate escalator statement?",
)
mock_load_field.return_value = escalator_statement_field
# Execute
result = dynamic_funcs.get_dynamic_one_to_one_fields(self.answer_dicts)
@patch('src.investment.dynamic_funcs.prompt_calls.prompt_dynamic_primary')
@patch("src.investment.dynamic_funcs.prompt_calls.prompt_dynamic_primary")
def test_dynamic_primary_with_header_answers(self, mock_prompt_dynamic_primary):
"""Test dynamic_primary function with answers in the header."""
# Setup
dynamic_fields = self.test_fieldset
reimbursement_level_fields = FieldSet()
exhibit_level_answers = {}
# Configure mocks
mock_prompt_dynamic_primary.side_effect = ["Single Answer", "Another Answer"]
# Execute
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic_primary(
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
self.constants, self.filename, dynamic_fields, reimbursement_level_fields
exhibit_level_answers, updated_reimbursement_fields = (
dynamic_funcs.dynamic_primary(
self.exhibit_text,
self.exhibit_header,
exhibit_level_answers,
self.constants,
self.filename,
dynamic_fields,
reimbursement_level_fields,
)
)
# Assert
self.assertEqual(len(exhibit_level_answers), 2) # Both fields should be added
self.assertEqual(exhibit_level_answers["TEST_FIELD"], "Single Answer")
@@ -217,110 +272,162 @@ class TestDynamicFuncsWithRealConstants(unittest.TestCase):
mock_prompt_dynamic_primary.assert_called()
self.assertEqual(mock_prompt_dynamic_primary.call_count, 2)
@patch('src.investment.dynamic_funcs.prompt_calls.prompt_dynamic_primary')
def test_dynamic_primary_with_multiple_header_answers(self, mock_prompt_dynamic_primary):
@patch("src.investment.dynamic_funcs.prompt_calls.prompt_dynamic_primary")
def test_dynamic_primary_with_multiple_header_answers(
self, mock_prompt_dynamic_primary
):
"""Test dynamic_primary function with multiple answers in the header."""
# Setup
dynamic_fields = self.test_fieldset
reimbursement_level_fields = FieldSet()
exhibit_level_answers = {}
# Configure mocks to return multiple values for first field
mock_prompt_dynamic_primary.side_effect = ["Value 1|Value 2|Value 3", "Another Answer"]
# Execute
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic_primary(
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
self.constants, self.filename, dynamic_fields, reimbursement_level_fields
)
# Assert
self.assertEqual(len(exhibit_level_answers), 1) # One field should be added to exhibit level
self.assertEqual(exhibit_level_answers["TEST_FIELD2"], "Another Answer")
self.assertEqual(len(updated_reimbursement_fields.fields), 1) # One field should be added to reimbursement level
self.assertEqual(updated_reimbursement_fields.fields[0].field_name, "TEST_FIELD")
self.assertEqual(updated_reimbursement_fields.fields[0].valid_values, "Value 1|Value 2|Value 3")
self.assertTrue("Ensure ALL values that apply" in updated_reimbursement_fields.fields[0].prompt)
@patch('src.investment.prompt_calls.prompt_dynamic')
# Configure mocks to return multiple values for first field
mock_prompt_dynamic_primary.side_effect = [
"Value 1|Value 2|Value 3",
"Another Answer",
]
# Execute
exhibit_level_answers, updated_reimbursement_fields = (
dynamic_funcs.dynamic_primary(
self.exhibit_text,
self.exhibit_header,
exhibit_level_answers,
self.constants,
self.filename,
dynamic_fields,
reimbursement_level_fields,
)
)
# Assert
self.assertEqual(
len(exhibit_level_answers), 1
) # One field should be added to exhibit level
self.assertEqual(exhibit_level_answers["TEST_FIELD2"], "Another Answer")
self.assertEqual(
len(updated_reimbursement_fields.fields), 1
) # One field should be added to reimbursement level
self.assertEqual(
updated_reimbursement_fields.fields[0].field_name, "TEST_FIELD"
)
self.assertEqual(
updated_reimbursement_fields.fields[0].valid_values,
"Value 1|Value 2|Value 3",
)
self.assertTrue(
"Ensure ALL values that apply"
in updated_reimbursement_fields.fields[0].prompt
)
@patch("src.investment.prompt_calls.prompt_dynamic")
def test_dynamic_with_real_constants(self, mock_prompt_dynamic):
"""Test dynamic function with real Constants object."""
# Setup
dynamic_fields = self.test_fieldset
reimbursement_level_fields = FieldSet()
exhibit_level_answers = {}
# Configure mocks
mock_prompt_dynamic.side_effect = [
{"TEST_FIELD": "Header Answer", "TEST_FIELD2": ""}, # Header results
{"TEST_FIELD2": "Text Answer"} # Text results (only for remaining field)
{"TEST_FIELD2": "Text Answer"}, # Text results (only for remaining field)
]
# Execute
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
self.constants, self.filename, dynamic_fields, reimbursement_level_fields
self.exhibit_text,
self.exhibit_header,
exhibit_level_answers,
self.constants,
self.filename,
dynamic_fields,
reimbursement_level_fields,
)
# Assert
self.assertEqual(len(updated_reimbursement_fields.fields), 2) # Both fields should be added to reimbursement level
self.assertEqual(
len(updated_reimbursement_fields.fields), 2
) # Both fields should be added to reimbursement level
# Check that the first field was processed from header results
self.assertEqual(updated_reimbursement_fields.fields[0].field_name, "TEST_FIELD")
self.assertEqual(updated_reimbursement_fields.fields[0].valid_values, "Header Answer")
self.assertEqual(
updated_reimbursement_fields.fields[0].field_name, "TEST_FIELD"
)
self.assertEqual(
updated_reimbursement_fields.fields[0].valid_values, "Header Answer"
)
# Check that the second field was processed from text results
self.assertEqual(updated_reimbursement_fields.fields[1].field_name, "TEST_FIELD2")
self.assertEqual(updated_reimbursement_fields.fields[1].valid_values, "Text Answer")
self.assertEqual(
updated_reimbursement_fields.fields[1].field_name, "TEST_FIELD2"
)
self.assertEqual(
updated_reimbursement_fields.fields[1].valid_values, "Text Answer"
)
# Verify that prompt_dynamic was called twice
self.assertEqual(mock_prompt_dynamic.call_count, 2)
@patch('src.investment.prompt_calls.prompt_dynamic')
@patch("src.investment.prompt_calls.prompt_dynamic")
def test_dynamic_with_no_fields(self, mock_prompt_dynamic):
"""Test dynamic function with an empty fieldset."""
# Setup
empty_fieldset = FieldSet()
reimbursement_level_fields = FieldSet()
exhibit_level_answers = {}
# Execute
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
self.constants, self.filename, empty_fieldset, reimbursement_level_fields
self.exhibit_text,
self.exhibit_header,
exhibit_level_answers,
self.constants,
self.filename,
empty_fieldset,
reimbursement_level_fields,
)
# Assert
self.assertEqual(len(exhibit_level_answers), 0)
self.assertEqual(len(updated_reimbursement_fields.fields), 0)
mock_prompt_dynamic.assert_not_called()
@patch('src.investment.prompt_calls.prompt_dynamic')
@patch("src.investment.prompt_calls.prompt_dynamic")
def test_dynamic_with_no_answers(self, mock_prompt_dynamic):
"""Test dynamic function when neither header nor text contains answers."""
# Setup
dynamic_fields = self.test_fieldset
reimbursement_level_fields = FieldSet()
exhibit_level_answers = {}
# Configure mocks - no answers anywhere
mock_prompt_dynamic.side_effect = [
{"TEST_FIELD": "", "TEST_FIELD2": ""}, # Header results
{"TEST_FIELD": "", "TEST_FIELD2": ""} # Text results
{"TEST_FIELD": "", "TEST_FIELD2": ""}, # Text results
]
# Execute
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
self.constants, self.filename, dynamic_fields, reimbursement_level_fields
self.exhibit_text,
self.exhibit_header,
exhibit_level_answers,
self.constants,
self.filename,
dynamic_fields,
reimbursement_level_fields,
)
# Assert
self.assertEqual(len(updated_reimbursement_fields.fields), 0) # No fields should be added to reimbursement level
self.assertEqual(exhibit_level_answers, {"TEST_FIELD": "", "TEST_FIELD2": ""}) # Both fields should be added to exhibit level
self.assertEqual(
len(updated_reimbursement_fields.fields), 0
) # No fields should be added to reimbursement level
self.assertEqual(
exhibit_level_answers, {"TEST_FIELD": "", "TEST_FIELD2": ""}
) # Both fields should be added to exhibit level
self.assertEqual(mock_prompt_dynamic.call_count, 2)
if __name__ == '__main__':
unittest.main()
if __name__ == "__main__":
unittest.main()
+3 -2
View File
@@ -1,10 +1,11 @@
import json
import os
import unittest
from unittest.mock import patch, mock_open
from unittest.mock import mock_open, patch
from src.prompts.fieldset import Field, FieldSet, _field_data_cache, _cache_lock
from constants.constants import Constants
from src.prompts.fieldset import (Field, FieldSet, _cache_lock,
_field_data_cache)
class TestFieldWithRealConstants(unittest.TestCase):
@@ -3,7 +3,6 @@ from unittest.mock import MagicMock, patch
import pandas as pd
import pytest
from src.investment.postprocessing_funcs import (
date_postprocess, deduplicate_provider_columns,
flatten_singleton_string_list, format_rate_fields_with_commas,
@@ -95,9 +94,7 @@ class TestPostprocessFunctions(unittest.TestCase):
@patch("src.utils.llm_utils.invoke_claude")
@patch("src.investment.postprocessing_funcs.invoke_derived_term_date")
@patch(
"src.prompts.fieldset.FieldSet.load_from_file"
) # Patch this specific method
@patch("src.prompts.fieldset.FieldSet.load_from_file") # Patch this specific method
def test_date_postprocess(
self, mock_load_from_file, mock_derived_term_date, mock_invoke_claude
):
+4 -11
View File
@@ -2,17 +2,10 @@ from io import StringIO
import pandas as pd
import pytest
from src.utils.io_utils import (
filter_already_processed,
list_s3_files,
read_input,
read_input_csv,
read_local,
read_s3,
read_s3_csv,
read_xlsb,
remove_txt_extension,
)
from src.utils.io_utils import (filter_already_processed, list_s3_files,
read_input, read_input_csv, read_local,
read_s3, read_s3_csv, read_xlsb,
remove_txt_extension)
class TestIOUtils:
File diff suppressed because it is too large Load Diff
@@ -1,5 +1,4 @@
import pytest
from src.investment.preprocessing_funcs import (clean_law_symbols,
filter_quick_review,
remove_page_indicators)
@@ -2,7 +2,6 @@ import unittest
from unittest.mock import patch
import pytest
from constants.constants import Constants
from src.investment.preprocess import clean_tables
@@ -1,5 +1,4 @@
import pytest
from src.investment.smart_chunking_funcs import parse_chunk, stitch_chunks
+5 -10
View File
@@ -1,19 +1,14 @@
import numpy as np
import pandas as pd
import pytest
import src.utils as utils
import src.utils.llm_utils as llm_utils
from constants.delimiters import Delimiter
from src.utils.string_utils import (
contains_reimbursement,
count_reimbursements_in_exhibit,
extract_signature_page,
extract_text_from_delimiters,
is_empty,
json_parsing_search,
page_key_sort,
)
from src.utils.string_utils import (contains_reimbursement,
count_reimbursements_in_exhibit,
extract_signature_page,
extract_text_from_delimiters, is_empty,
json_parsing_search, page_key_sort)
class TestStringUtils:
@@ -3,7 +3,6 @@ import unittest
from unittest.mock import patch
import pandas as pd
from constants.constants import Constants
from src.investment.table_funcs import (END_MARKER, START_MARKER, Table,
combine_continuous_tables,
+1 -2
View File
@@ -2,9 +2,8 @@ import json
import re
from unittest.mock import MagicMock, patch
import pytest
import constants.regex_patterns as regex_patterns
import pytest
import src.investment.tin_npi_funcs as tin_npi_funcs
from src.prompts.fieldset import FieldSet