doczyai-pipelines/fieldExtraction/tests/test_document_classification.py

"""
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()