""" Unit tests for document_classification/main.py module. Tests cover: - call_llm: LLM invocation wrapper - save_result_to_json: JSON file saving - process_single_file: Single document classification - main: Full batch classification with S3/local upload """ import json import os import tempfile import unittest from datetime import datetime from unittest.mock import MagicMock, Mock, patch, call import pandas as pd from src.document_classification import main as dtc_main from src.utils import io_utils class TestCallLLM(unittest.TestCase): """Test the call_llm function that wraps llm_utils.invoke_claude.""" @patch('src.document_classification.main.llm_utils.invoke_claude') def test_call_llm_success(self, mock_invoke): """Test successful LLM call returns cleaned response.""" mock_invoke.return_value = " YES \n" result = dtc_main.call_llm( filename="test.txt", context="This is a Service Agreement...", question="Is this a contract?" ) self.assertEqual(result, "YES") mock_invoke.assert_called_once() call_args = mock_invoke.call_args self.assertIn("This is a Service Agreement", call_args.kwargs['prompt']) self.assertEqual(call_args.kwargs['model_id'], "sonnet_latest") self.assertEqual(call_args.kwargs['filename'], "test.txt") self.assertEqual(call_args.kwargs['max_tokens'], 8192) @patch('src.document_classification.main.llm_utils.invoke_claude') @patch('src.document_classification.main.logging') def test_call_llm_error_handling(self, mock_logging, mock_invoke): """Test error handling returns ERROR and logs the exception.""" mock_invoke.side_effect = Exception("API timeout") result = dtc_main.call_llm( filename="test.txt", context="Document text", question="Question" ) self.assertEqual(result, "ERROR") mock_logging.error.assert_called_once() self.assertIn("test.txt", mock_logging.error.call_args[0][0]) @patch('src.document_classification.main.llm_utils.invoke_claude') def test_call_llm_prompt_format(self, mock_invoke): """Test that prompt is correctly formatted with context and question.""" mock_invoke.return_value = "Answer" context = "Contract document text here" question = "What type is this?" dtc_main.call_llm("test.txt", context, question) prompt = mock_invoke.call_args.kwargs['prompt'] self.assertIn("Document text:", prompt) self.assertIn(context, prompt) self.assertIn(question, prompt) class TestSaveResultToJson(unittest.TestCase): """Test the io_utils.save_result_to_json function for JSON file operations.""" def setUp(self): """Create a temporary directory for test files.""" self.temp_dir = tempfile.mkdtemp() def tearDown(self): """Clean up temporary directory.""" import shutil shutil.rmtree(self.temp_dir, ignore_errors=True) def test_save_result_creates_folder(self): """Test that function creates output folder if it doesn't exist.""" json_folder = os.path.join(self.temp_dir, "new_folder") self.assertFalse(os.path.exists(json_folder)) io_utils.save_result_to_json( filename="test.txt", result={"is_contract": "YES"}, json_folder=json_folder ) self.assertTrue(os.path.exists(json_folder)) def test_save_result_creates_json_file(self): """Test that JSON file is created with correct content.""" result = { "is_contract": "YES", "contract_type": "Service Agreement", "found_on_page": 1 } io_utils.save_result_to_json( filename="contract123.txt", result=result, json_folder=self.temp_dir ) json_file = os.path.join(self.temp_dir, "contract123.txt.json") self.assertTrue(os.path.exists(json_file)) with open(json_file, 'r') as f: saved_data = json.load(f) self.assertEqual(saved_data['filename'], "contract123.txt") self.assertEqual(saved_data['result'], result) def test_save_result_sanitizes_filename(self): """Test that invalid filename characters are replaced.""" result = {"is_contract": "NO"} io_utils.save_result_to_json( filename="path/to/file:with:colons.txt", result=result, json_folder=self.temp_dir ) # Should replace / : \ with _ expected_filename = "path_to_file_with_colons.txt.json" json_file = os.path.join(self.temp_dir, expected_filename) self.assertTrue(os.path.exists(json_file)) def test_save_result_adds_json_extension(self): """Test that .json extension is added if not present.""" io_utils.save_result_to_json( filename="test.txt", result={"is_contract": "YES"}, json_folder=self.temp_dir ) json_file = os.path.join(self.temp_dir, "test.txt.json") self.assertTrue(os.path.exists(json_file)) class TestProcessSingleFile(unittest.TestCase): """Test the process_single_file function for document classification.""" def setUp(self): """Create temporary directory for JSON output.""" self.temp_dir = tempfile.mkdtemp() def tearDown(self): """Clean up temporary directory.""" import shutil shutil.rmtree(self.temp_dir, ignore_errors=True) @patch('src.document_classification.main.call_llm') @patch('src.document_classification.main.split_text') @patch('src.document_classification.main.clean_law_symbols') @patch('src.document_classification.main.clean_newlines') def test_process_contract_found_first_page(self, mock_clean_nl, mock_clean_law, mock_split, mock_llm): """Test processing when contract is found on first page.""" mock_clean_nl.return_value = "cleaned text" mock_clean_law.return_value = "cleaned text" mock_split.return_value = { "1": "This is a Service Agreement between parties...", "2": "Page 2 content" } # First call: is_contract check returns YES # Second call: contract_type returns Service Agreement mock_llm.side_effect = ["YES", "Service Agreement"] filename, result = dtc_main.process_single_file( filename="contract.txt", contract_text="raw contract text", json_folder=self.temp_dir ) self.assertEqual(filename, "contract.txt") self.assertEqual(result['is_contract'], "YES") self.assertEqual(result['contract_type'], "Service Agreement") self.assertEqual(result['non_contract_type'], "N/A") self.assertEqual(result['found_on_page'], 1) # Should make 2 LLM calls (is_contract + contract_type) self.assertEqual(mock_llm.call_count, 2) @patch('src.document_classification.main.call_llm') @patch('src.document_classification.main.split_text') @patch('src.document_classification.main.clean_law_symbols') @patch('src.document_classification.main.clean_newlines') def test_process_non_contract(self, mock_clean_nl, mock_clean_law, mock_split, mock_llm): """Test processing when document is not a contract.""" mock_clean_nl.return_value = "cleaned text" mock_clean_law.return_value = "cleaned text" mock_split.return_value = { "1": "This is an email from John...", "2": "Email continues..." } # First call: is_contract returns NO # Second call: non_contract_type returns Email mock_llm.side_effect = ["NO", "NO", "Email"] filename, result = dtc_main.process_single_file( filename="email.txt", contract_text="email text", json_folder=self.temp_dir ) self.assertEqual(result['is_contract'], "NO") self.assertEqual(result['contract_type'], "N/A") self.assertEqual(result['non_contract_type'], "Email") self.assertEqual(result['found_on_page'], 1) @patch('src.document_classification.main.call_llm') @patch('src.document_classification.main.split_text') @patch('src.document_classification.main.clean_law_symbols') @patch('src.document_classification.main.clean_newlines') @patch('src.document_classification.main.config.DTC_MAX_PAGES_TO_CHECK', 3) def test_process_contract_found_later_page(self, mock_clean_nl, mock_clean_law, mock_split, mock_llm): """Test when contract signature is found on page 3.""" mock_clean_nl.return_value = "cleaned" mock_clean_law.return_value = "cleaned" mock_split.return_value = { "1": "Cover page", "2": "Table of contents", "3": "SERVICE AGREEMENT between parties", "4": "More content" } # Page 1: NO, Page 2: NO, Page 3: YES, then get type mock_llm.side_effect = ["NO", "NO", "YES", "Service Agreement"] filename, result = dtc_main.process_single_file( filename="contract.txt", contract_text="text", json_folder=self.temp_dir ) self.assertEqual(result['is_contract'], "YES") self.assertEqual(result['found_on_page'], 3) @patch('src.document_classification.main.call_llm') @patch('src.document_classification.main.split_text') @patch('src.document_classification.main.clean_law_symbols') @patch('src.document_classification.main.clean_newlines') @patch('src.document_classification.main.logging') def test_process_error_handling(self, mock_logging, mock_clean_nl, mock_clean_law, mock_split, mock_llm): """Test error handling returns error result.""" mock_clean_nl.side_effect = Exception("Processing error") filename, result = dtc_main.process_single_file( filename="bad.txt", contract_text="text", json_folder=self.temp_dir ) self.assertEqual(result['is_contract'], "ERROR") self.assertEqual(result['contract_type'], "ERROR") self.assertEqual(result['non_contract_type'], "ERROR") self.assertEqual(result['found_on_page'], "ERROR") mock_logging.error.assert_called_once() @patch('src.document_classification.main.call_llm') @patch('src.document_classification.main.split_text') @patch('src.document_classification.main.clean_law_symbols') @patch('src.document_classification.main.clean_newlines') @patch('src.document_classification.main.config.DTC_MAX_PAGES_TO_CHECK', 2) def test_process_handles_fewer_pages(self, mock_clean_nl, mock_clean_law, mock_split, mock_llm): """Test when document has fewer pages than MAX_PAGES_TO_CHECK.""" mock_clean_nl.return_value = "cleaned" mock_clean_law.return_value = "cleaned" mock_split.return_value = {"1": "Only one page"} # Only 1 page mock_llm.side_effect = ["NO", "Memo"] filename, result = dtc_main.process_single_file( filename="short.txt", contract_text="text", json_folder=self.temp_dir ) # Should only check page 1, not error on page 2 self.assertEqual(result['is_contract'], "NO") class TestRunDTCClassification(unittest.TestCase): """Integration tests for main function.""" def setUp(self): """Create temporary directory for outputs.""" self.temp_dir = tempfile.mkdtemp() def tearDown(self): """Clean up temporary directory.""" import shutil shutil.rmtree(self.temp_dir, ignore_errors=True) @patch('src.document_classification.main.io_utils.write_s3') @patch('src.document_classification.main.io_utils.write_local') @patch('src.document_classification.main.process_single_file') @patch('src.document_classification.main.config.WRITE_TO_S3', False) @patch('src.document_classification.main.config.BATCH_ID', 'test_batch') @patch('src.document_classification.main.logging') def test_main_local(self, mock_logging, mock_process, mock_write_local, mock_write_s3): """Test full DTC run with local file output.""" # Mock process_single_file to return results mock_process.side_effect = [ ("file1.txt", {"is_contract": "YES", "contract_type": "Service Agreement", "non_contract_type": "N/A", "found_on_page": 1}), ("file2.txt", {"is_contract": "NO", "contract_type": "N/A", "non_contract_type": "Email", "found_on_page": 1}) ] input_dict = { "file1.txt": "Service Agreement text...", "file2.txt": "Email text..." } run_timestamp = "run_20250106_10-30_test" results = dtc_main.main(input_dict, run_timestamp) # Check results returned correctly self.assertIn("file1.txt", results) self.assertIn("file2.txt", results) self.assertEqual(results["file1.txt"]["is_contract"], "YES") self.assertEqual(results["file2.txt"]["is_contract"], "NO") # Check write_local was called twice (once for details, once for summary), not write_s3 self.assertEqual(mock_write_local.call_count, 2) mock_write_s3.assert_not_called() # Verify write_local called with correct args for main results first_call_args = mock_write_local.call_args_list[0] df_arg = first_call_args[0][0] self.assertIsInstance(df_arg, pd.DataFrame) self.assertEqual(len(df_arg), 2) self.assertEqual(first_call_args[0][2], run_timestamp) self.assertEqual(first_call_args[0][3], "dtc") # Verify summary was also saved second_call_args = mock_write_local.call_args_list[1] summary_df = second_call_args[0][0] self.assertIsInstance(summary_df, pd.DataFrame) self.assertEqual(second_call_args[0][3], "dtc_summary") @patch('src.document_classification.main.io_utils.write_s3') @patch('src.document_classification.main.io_utils.write_local') @patch('src.document_classification.main.process_single_file') @patch('src.document_classification.main.config.WRITE_TO_S3', True) @patch('src.document_classification.main.config.BATCH_ID', 'test_batch') @patch('src.document_classification.main.logging') def test_main_s3(self, mock_logging, mock_process, mock_write_local, mock_write_s3): """Test full DTC run with S3 upload.""" mock_process.return_value = ( "file1.txt", {"is_contract": "YES", "contract_type": "MSA", "non_contract_type": "N/A", "found_on_page": 2} ) input_dict = {"file1.txt": "Master Service Agreement..."} run_timestamp = "run_20250106_10-30_test" results = dtc_main.main(input_dict, run_timestamp) # Check S3 upload was called twice (main + summary), not local self.assertEqual(mock_write_s3.call_count, 2) mock_write_local.assert_not_called() # Verify write_s3 called with correct args for main results first_call = mock_write_s3.call_args_list[0] self.assertEqual(first_call[0][2], run_timestamp) self.assertEqual(first_call[0][3], "dtc") # Verify summary was also saved to S3 second_call = mock_write_s3.call_args_list[1] self.assertEqual(second_call[0][2], run_timestamp) self.assertEqual(second_call[0][3], "dtc_summary") @patch('src.document_classification.main.io_utils.write_local') @patch('src.document_classification.main.process_single_file') @patch('src.document_classification.main.config.WRITE_TO_S3', False) @patch('src.document_classification.main.config.BATCH_ID', 'test_batch') @patch('src.document_classification.main.logging') def test_run_dtc_summary_statistics(self, mock_logging, mock_process, mock_write_local): """Test that summary statistics are logged correctly.""" # 2 contracts, 1 non-contract, 1 error mock_process.side_effect = [ ("f1.txt", {"is_contract": "YES", "contract_type": "Service Agreement", "non_contract_type": "N/A", "found_on_page": 1}), ("f2.txt", {"is_contract": "YES", "contract_type": "MSA", "non_contract_type": "N/A", "found_on_page": 1}), ("f3.txt", {"is_contract": "NO", "contract_type": "N/A", "non_contract_type": "Email", "found_on_page": 1}), ("f4.txt", {"is_contract": "ERROR", "contract_type": "ERROR", "non_contract_type": "ERROR", "found_on_page": "ERROR"}) ] input_dict = {f"f{i}.txt": "text" for i in range(1, 5)} dtc_main.main(input_dict, "run_test") # Check that summary was logged log_calls = [str(call) for call in mock_logging.info.call_args_list] summary_log = [c for c in log_calls if "DTC Summary" in c] self.assertTrue(len(summary_log) > 0) # Verify counts in summary summary_str = str(summary_log[0]) self.assertIn("2 contracts", summary_str) self.assertIn("1 non-contracts", summary_str) self.assertIn("1 errors", summary_str) @patch('src.document_classification.main.io_utils.write_local') @patch('src.document_classification.main.process_single_file') @patch('src.document_classification.main.config.WRITE_TO_S3', False) @patch('src.document_classification.main.config.BATCH_ID', 'test_batch') @patch('src.document_classification.main.logging') def test_run_dtc_no_contracts_summary(self, mock_logging, mock_process, mock_write_local): """Test summary when there are no contracts (covers else block).""" # All non-contracts mock_process.side_effect = [ ("f1.txt", {"is_contract": "NO", "contract_type": "N/A", "non_contract_type": "Email", "found_on_page": 1}), ("f2.txt", {"is_contract": "NO", "contract_type": "N/A", "non_contract_type": "Form", "found_on_page": 1}), ] input_dict = {f"f{i}.txt": "text" for i in range(1, 3)} dtc_main.main(input_dict, "run_test") # Verify write_local was called twice self.assertEqual(mock_write_local.call_count, 2) # Check summary DataFrame summary_df = mock_write_local.call_args_list[1][0][0] self.assertEqual(summary_df['contracts'].iloc[0], 0) self.assertEqual(summary_df['non_contracts'].iloc[0], 2) self.assertEqual(summary_df['most_common_contract_type'].iloc[0], "N/A") @patch('src.document_classification.main.io_utils.write_local') @patch('src.document_classification.main.process_single_file') @patch('src.document_classification.main.config.WRITE_TO_S3', False) @patch('src.document_classification.main.config.BATCH_ID', 'test_batch') @patch('src.document_classification.main.logging') def test_run_dtc_future_exception(self, mock_logging, mock_process, mock_write_local): """Test exception handling in futures loop (covers lines 244-246).""" # Simulate an exception being raised by a future def raise_exception(filename, text, folder): if filename == "bad.txt": raise ValueError("Simulated processing error") return (filename, {"is_contract": "YES", "contract_type": "MSA", "non_contract_type": "N/A", "found_on_page": 1}) mock_process.side_effect = raise_exception input_dict = { "good.txt": "text", "bad.txt": "text" } results = dtc_main.main(input_dict, "run_test") # Should handle exception and still return results self.assertIn("bad.txt", results) self.assertEqual(results["bad.txt"]["is_contract"], "ERROR") # Verify error was logged error_logs = [call for call in mock_logging.error.call_args_list if "Exception occurred" in str(call)] self.assertTrue(len(error_logs) > 0) @patch('src.document_classification.main.call_llm') @patch('src.document_classification.main.split_text') @patch('src.document_classification.main.clean_law_symbols') @patch('src.document_classification.main.clean_newlines') @patch('src.document_classification.main.config.DTC_MAX_PAGES_TO_CHECK', 10) def test_process_progress_logging_every_5(self, mock_clean_nl, mock_clean_law, mock_split, mock_llm): """Test that progress is logged every 5 files (covers line 169).""" temp_dir = tempfile.mkdtemp() try: mock_clean_nl.return_value = "cleaned" mock_clean_law.return_value = "cleaned" mock_split.return_value = {"1": "Contract text"} mock_llm.side_effect = ["YES", "Service Agreement"] # Set up progress counter to trigger line 169 from src.document_classification.main import progress_counter, progress_lock with progress_lock: progress_counter["total"] = 10 progress_counter["completed"] = 4 # Next will be 5 filename, result = dtc_main.process_single_file( filename="file5.txt", contract_text="text", json_folder=temp_dir ) # Verify file was processed self.assertEqual(filename, "file5.txt") self.assertEqual(result["is_contract"], "YES") finally: import shutil shutil.rmtree(temp_dir, ignore_errors=True) if __name__ == '__main__': unittest.main()