doczyai-pipelines/fieldExtraction/tests/test_io_utils.py

from io import StringIO
import os

import pandas as pd
import pytest
from botocore.exceptions import ClientError
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, write_s3)


class TestIOUtils:
    @pytest.fixture(autouse=True)
    def cleanup_config_patches(self, mocker):
        """Ensure config patches are cleaned up after each test to prevent leakage."""
        # Store original config values before test
        import src.config as config_module
        original_values = {}
        config_attrs = [
            "S3_CLIENT",
            "BATCH_ID",
            "S3_OUTPUT_BUCKET",
            "READ_MODE",
            "DF_READ_MODE",
            "AC_DF",
            "B_DF",
            "WRITE_TO_S3",
            "OUTPUT_DIRECTORY",
            "AC_RESULTS_NAME",
            "B_RESULTS_NAME",
        ]
        for attr in config_attrs:
            if hasattr(config_module, attr):
                original_values[attr] = getattr(config_module, attr)

        yield  # Run the test

        # Restore original values after test
        for attr, value in original_values.items():
            setattr(config_module, attr, value)
    # Fixtures for reusable test data
    @pytest.fixture
    def mock_csv_data(self):
        return "col1,col2\n1,2\n3,4"

    @pytest.fixture
    def mock_txt_data(self):
        return "This is a test file."

    @pytest.fixture
    def mock_s3_response(self):
        return {
            "Contents": [
                {"Key": "file1.txt"},
                {"Key": "file2.csv"},
                {"Key": "file3.xlsx"},
            ]
        }

    @pytest.fixture
    def mock_s3_csv_response(self):
        return {"Body": StringIO("col1,col2\n1,2\n3,4")}

    # Tests for read_local
    def test_read_local_txt_utf8(self, mock_txt_data, mocker):
        mocker.patch("builtins.open", mocker.mock_open(read_data=mock_txt_data))
        mocker.patch("os.path.isfile", return_value=True)
        result = read_local("test.txt")
        assert result == mock_txt_data

    def test_read_local_csv(self, mock_csv_data, mocker):
        mock_df = pd.read_csv(StringIO(mock_csv_data))
        mocker.patch("pandas.read_csv", return_value=mock_df)
        mocker.patch("os.path.isfile", return_value=True)
        result = read_local("test.csv")
        pd.testing.assert_frame_equal(result, mock_df)

    def test_read_local_nonexistent_file(self, mocker):
        mocker.patch("os.path.isfile", return_value=False)
        result = read_local("nonexistent.txt")
        assert result is None

    def test_list_s3_files(self, mocker):
        # Test case 1: Empty S3 bucket
        mocker.patch(
            "src.config.S3_CLIENT.list_objects_v2", return_value={"Contents": []}
        )
        result = list_s3_files()
        assert result == []

        # Test case 2: S3 bucket with dummy files
        mock_s3_response = {
            "Contents": [
                {"Key": "file1.txt"},
                {"Key": "file2.csv"},
                {"Key": "file3.xlsx"},
            ]
        }
        mocker.patch(
            "src.config.S3_CLIENT.list_objects_v2", return_value=mock_s3_response
        )
        result = list_s3_files()
        assert result == ["file1.txt", "file2.csv", "file3.xlsx"]

    # Tests for read_s3
    def test_read_s3(self, mock_txt_data, mocker):
        mock_s3_client = mocker.MagicMock()
        mock_s3_client.list_objects_v2.return_value = {
            "Contents": [{"Key": "file1.txt"}],
            "IsTruncated": False,
        }
        mock_s3_client.get_object.return_value = {
            "Body": mocker.MagicMock(
                read=mocker.MagicMock(return_value=mock_txt_data.encode("utf-8"))
            )
        }
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        result = read_s3()
        assert result == {"file1.txt": mock_txt_data}

    # Tests for read_s3_csv
    def test_read_s3_csv(self, mock_csv_data, mocker):
        mock_s3_client = mocker.MagicMock()
        mock_s3_client.list_objects_v2.return_value = {
            "Contents": [{"Key": "file2.csv"}],
        }
        # Mock get_object to return a response with a bytes object
        mock_s3_client.get_object.return_value = {
            "Body": mocker.MagicMock(
                read=mocker.MagicMock(return_value=mock_csv_data.encode("utf-8"))
            )
        }
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        result = read_s3_csv("s3://bucket/file2.csv")
        expected_df = pd.read_csv(StringIO(mock_csv_data))
        pd.testing.assert_frame_equal(result, expected_df)

        # Verify the S3 client was called correctly
        mock_s3_client.list_objects_v2.assert_called_once_with(
            Bucket="bucket", Prefix="file2.csv"
        )
        mock_s3_client.get_object.assert_called_once_with(
            Bucket="bucket", Key="file2.csv"
        )

    # Tests for read_input
    def test_read_input_local(self, mock_txt_data, mocker):
        mocker.patch("os.listdir", return_value=["file1.txt"])
        mocker.patch("os.path.isfile", return_value=True)
        mocker.patch("src.utils.io_utils.read_local", return_value=mock_txt_data)
        mocker.patch("src.config.READ_MODE", "local")
        result = read_input()
        assert result == {"file1.txt": mock_txt_data}

    def test_read_input_s3(self, mock_txt_data, mocker):
        mocker.patch(
            "src.utils.io_utils.read_s3", return_value={"file1.txt": mock_txt_data}
        )
        mocker.patch("src.config.READ_MODE", "s3")
        result = read_input()
        assert result == {"file1.txt": mock_txt_data}

    # Tests for read_input_csv
    def test_read_input_csv_local(self, mock_csv_data, mocker):
        mock_df = pd.read_csv(StringIO(mock_csv_data))
        mock_read_local = mocker.patch(
            "src.utils.io_utils.read_local", return_value=mock_df
        )
        mocker.patch("src.config.DF_READ_MODE", "local")
        mocker.patch("src.config.AC_DF", "ac.csv")
        mocker.patch("src.config.B_DF", "b.csv")

        ac_df, b_df = read_input_csv()

        pd.testing.assert_frame_equal(ac_df, mock_df)
        pd.testing.assert_frame_equal(b_df, mock_df)

        # Verify read_local was called with the correct paths
        mock_read_local.assert_any_call("ac.csv")
        mock_read_local.assert_any_call("b.csv")

    def test_read_input_csv_s3(self, mock_csv_data, mocker):
        # Create a mock DataFrame
        mock_df = pd.read_csv(StringIO(mock_csv_data))

        # Mock read_s3_csv to return the mock DataFrame
        mock_read_s3_csv = mocker.patch(
            "src.utils.io_utils.read_s3_csv", return_value=mock_df
        )

        # Set the src.config to use S3 mode
        mocker.patch("src.config.DF_READ_MODE", "s3")
        mocker.patch("src.config.AC_DF", "s3://bucket/ac.csv")
        mocker.patch("src.config.B_DF", "s3://bucket/b.csv")

        # Call the function under test
        ac_df, b_df = read_input_csv()

        # Verify the results
        pd.testing.assert_frame_equal(ac_df, mock_df)
        pd.testing.assert_frame_equal(b_df, mock_df)

        # Verify read_s3_csv was called with the correct S3 paths
        mock_read_s3_csv.assert_any_call("s3://bucket/ac.csv")
        mock_read_s3_csv.assert_any_call("s3://bucket/b.csv")

    def test_filter_already_processed_s3(self, mocker):
        # Case 1: WRITE_TO_S3 is True (S3 mode)
        # Mock src.config.BATCH_ID
        mock_batch_id = "mock_batch_123"
        mocker.patch("src.config.BATCH_ID", mock_batch_id)
        mock_get_processed_files = mocker.patch(
            "src.tracking.tracking_utils.get_processed_files",
            return_value=(set(["file1.txt"]), set([])),
        )
        mock_write_to_s3 = mocker.patch("src.config.WRITE_TO_S3", True)
        input_dict = {"file1.txt": "content1", "file2.txt": "content2"}
        input_dict_b = {"file1.txt": "content1", "file2.txt": "content2"}
        result_ac, result_b = filter_already_processed(input_dict, input_dict_b)
        assert result_ac == {"file2.txt": "content2"}
        assert result_b == {"file1.txt": "content1", "file2.txt": "content2"}

        # Verify src.tracking.tracking_utils.get_processed_files was called
        mock_get_processed_files.assert_called_once_with(mock_batch_id)

    def test_filter_already_processed_local(self, mocker):
        # Case 2: WRITE_TO_S3 is False (local mode)
        mocker.patch("src.config.WRITE_TO_S3", False)
        mocker.patch("src.config.OUTPUT_DIRECTORY", "/mock/output/dir")
        mocker.patch("src.config.AC_RESULTS_NAME", "ac_results.txt")
        mocker.patch("src.config.B_RESULTS_NAME", "b_results.txt")

        # Mock os.listdir to return a list of folders for the top-level directory
        mocker.patch(
            "os.listdir",
            side_effect=lambda x: {
                "/mock/output/dir": ["folder1", "folder2"],  # Top-level directory
                "/mock/output/dir/folder1": ["ac_results.txt"],  # Folder 1 contents
                "/mock/output/dir/folder2": ["b_results.txt"],  # Folder 2 contents
            }[x],
        )

        # Mock os.path.isdir to return True for folders
        mocker.patch("os.path.isdir", return_value=True)

        # Mock os.path.join to construct paths
        mocker.patch("os.path.join", side_effect=lambda *args: "/".join(args))

        # Input dictionaries
        input_dict = {
            "folder1.txt": "content1",
            "folder2.txt": "content2",
            "folder3.txt": "content3",
        }
        input_dict_b = {
            "folder1.txt": "content1",
            "folder2.txt": "content2",
            "folder3.txt": "content3",
        }

        # Call the function under test
        result_ac, result_b = filter_already_processed(input_dict, input_dict_b)

        # Verify the results
        assert result_ac == {
            "folder2.txt": "content2",
            "folder3.txt": "content3",
        }  # folder1.txt is already processed
        assert result_b == {
            "folder1.txt": "content1",
            "folder3.txt": "content3",
        }  # folder2.txt is already processed

    # Tests for remove_txt_extension
    def test_remove_txt_extension(self):
        assert remove_txt_extension("file.txt") == "file"
        assert remove_txt_extension("file.csv") == "file.csv"
        assert remove_txt_extension(123) == 123
        assert remove_txt_extension(None) == None

    def test_read_xlsb(self, mocker):
        # Mock the sheet
        mock_sheet = mocker.MagicMock()

        # Mock the rows returned by the sheet
        # Use a list of rows, where each row is a list of MagicMock objects with a 'v' attribute
        mock_rows = [
            [
                mocker.MagicMock(v="Header1"),
                mocker.MagicMock(v="Header2"),
            ],  # Header row
            [mocker.MagicMock(v="Data1"), mocker.MagicMock(v="Data2")],  # Data row 1
            [mocker.MagicMock(v="Data3"), mocker.MagicMock(v="Data4")],  # Data row 2
        ]

        # Mock sheet.rows to return an iterable (e.g., a generator)
        mock_sheet.rows.return_value = iter(
            mock_rows
        )  # Use return_value for the method

        # Mock the workbook
        mock_workbook = mocker.MagicMock()
        # Ensure get_sheet() returns the mocked sheet
        mock_workbook.get_sheet.return_value.__enter__.return_value = mock_sheet

        # Mock open_workbook to return the mocked workbook
        mock_open_workbook = mocker.patch("src.utils.io_utils.open_workbook")
        mock_open_workbook.return_value.__enter__.return_value = mock_workbook

        # Call the function
        result = read_xlsb("dummy_path.xlsb")

        # Expected DataFrame
        expected_df = pd.DataFrame(
            [
                ["Data1", "Data2"],  # Data row 1
                ["Data3", "Data4"],  # Data row 2
            ],
            columns=["Header1", "Header2"],  # Header row
        )

        # Assert that the result matches the expected DataFrame
        pd.testing.assert_frame_equal(result, expected_df)

        # Verify that the mock was called correctly
        mock_open_workbook.assert_called_once_with("dummy_path.xlsb")
        mock_workbook.get_sheet.assert_called_once()

    # Tests for write_s3
    def test_write_s3_final(self, mocker):
        """Test write_s3 with final output_type as used in main.py."""
        mock_s3_client = mocker.MagicMock()
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        mocker.patch("src.config.BATCH_ID", "test_batch_123")
        mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket")
        mocker.patch("src.utils.io_utils.logging")

        df = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
        write_s3(df, "", "20240101120000", "final")

        mock_s3_client.put_object.assert_called_once()
        call_kwargs = mock_s3_client.put_object.call_args[1]
        assert call_kwargs["Bucket"] == "test-bucket"
        assert call_kwargs["Key"] == "test_batch_123/20240101120000/test_batch_123-RESULTS.csv"
        assert isinstance(call_kwargs["Body"], bytes)
        assert b"col1,col2" in call_kwargs["Body"]

    def test_write_s3_error(self, mocker):
        """Test write_s3 with error output_type as used in main.py."""
        mock_s3_client = mocker.MagicMock()
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        mocker.patch("src.config.BATCH_ID", "test_batch_123")
        mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket")
        mocker.patch("src.utils.io_utils.logging")

        df = pd.DataFrame({"error": ["test error"]})
        write_s3(df, "", "20240101120000", "error")

        call_kwargs = mock_s3_client.put_object.call_args[1]
        assert call_kwargs["Key"] == "test_batch_123/20240101120000/test_batch_123-ERRORS.csv"

    def test_write_s3_usage(self, mocker):
        """Test write_s3 with usage output_type as used in main.py."""
        mock_s3_client = mocker.MagicMock()
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        mocker.patch("src.config.BATCH_ID", "test_batch_123")
        mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket")
        mocker.patch("src.utils.io_utils.logging")

        df = pd.DataFrame({"filename": ["file1.txt"], "tokens": [100]})
        write_s3(df, "", "20240101120000", "usage")

        call_kwargs = mock_s3_client.put_object.call_args[1]
        assert call_kwargs["Key"] == "test_batch_123/20240101120000/tracking/test_batch_123-USAGE.csv"

    def test_write_s3_usage_summary(self, mocker):
        """Test write_s3 with usage_summary output_type as used in main.py."""
        mock_s3_client = mocker.MagicMock()
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        mocker.patch("src.config.BATCH_ID", "test_batch_123")
        mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket")
        mocker.patch("src.utils.io_utils.logging")

        df = pd.DataFrame({"total_tokens": [1000], "total_cost": [0.50]})
        write_s3(df, "", "20240101120000", "usage_summary")

        call_kwargs = mock_s3_client.put_object.call_args[1]
        assert call_kwargs["Key"] == "test_batch_123/20240101120000/tracking/test_batch_123-USAGE-SUMMARY.csv"

    def test_write_s3_client_error(self, mocker):
        """Test write_s3 error handling when S3 upload fails."""
        mock_s3_client = mocker.MagicMock()
        mock_s3_client.put_object.side_effect = ClientError(
            {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}},
            "PutObject",
        )
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        mocker.patch("src.config.BATCH_ID", "test_batch_123")
        mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket")
        mock_logger = mocker.patch("src.utils.io_utils.logging")

        df = pd.DataFrame({"col1": [1]})
        result = write_s3(df, "", "20240101120000", "final")

        assert result is None
        mock_logger.error.assert_called_once()
        assert "Error uploading to S3" in str(mock_logger.error.call_args)

    def test_write_s3_qc_qa_validated(self, mocker):
        """Test write_s3 with qc_qa_validated output type."""
        from src.utils.io_utils import write_local

        mock_s3_client = mocker.MagicMock()
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        mocker.patch("src.config.BATCH_ID", "test_batch_qc")
        mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket")

        df = pd.DataFrame({"FILE_NAME": ["test.pdf"], "QC_FLAG": ["Valid"]})
        write_s3(df, "", "run_20241114_10-00_test", "qc_qa_validated")

        # Verify S3 client was called
        mock_s3_client.put_object.assert_called_once()
        call_args = mock_s3_client.put_object.call_args

        # Check bucket
        assert call_args[1]["Bucket"] == "test-bucket"

        # Check key includes automation_qa-qc folder
        assert "automation_qa-qc" in call_args[1]["Key"]
        assert "QC-QA-VALIDATED.csv" in call_args[1]["Key"]

    def test_write_s3_qc_qa_stats(self, mocker):
        """Test write_s3 with qc_qa_stats output type."""
        mock_s3_client = mocker.MagicMock()
        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
        mocker.patch("src.config.BATCH_ID", "test_batch_qc")
        mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket")

        df = pd.DataFrame({"Field": ["RATE"], "Percentage_Filled": [100.0]})
        write_s3(df, "", "run_20241114_10-00_test", "qc_qa_stats")

        # Verify S3 client was called
        mock_s3_client.put_object.assert_called_once()
        call_args = mock_s3_client.put_object.call_args

        # Check bucket
        assert call_args[1]["Bucket"] == "test-bucket"

        # Check key includes automation_qa-qc folder and correct filename
        assert "automation_qa-qc" in call_args[1]["Key"]
        assert "QC-QA-STATS.csv" in call_args[1]["Key"]

    def test_write_local_qc_qa_validated(self, mocker, tmp_path):
        """Test write_local with qc_qa_validated output type."""
        from src.utils.io_utils import write_local

        mocker.patch("src.config.BATCH_ID", "test_batch_qc")
        # Mock os.makedirs and os.path.join to use tmp_path
        original_makedirs = os.makedirs
        original_join = os.path.join

        def mock_makedirs(path, exist_ok=False):
            # Use tmp_path for test isolation
            test_path = tmp_path / path
            test_path.mkdir(parents=True, exist_ok=exist_ok)
            return str(test_path)

        mocker.patch("os.makedirs", side_effect=mock_makedirs)
        mock_to_csv = mocker.patch.object(pd.DataFrame, "to_csv")

        df = pd.DataFrame({"FILE_NAME": ["test.pdf"], "QC_FLAG": ["Valid"]})
        write_local(df, "", "run_20241114_10-00_test", "qc_qa_validated")

        # Verify to_csv was called
        mock_to_csv.assert_called_once()
        call_args = mock_to_csv.call_args[0][0]

        # Check output path includes qa_qc_output folder
        assert "qa_qc_output" in call_args
        assert "QC-QA-VALIDATED.csv" in call_args

    def test_write_local_qc_qa_stats(self, mocker, tmp_path):
        """Test write_local with qc_qa_stats output type."""
        from src.utils.io_utils import write_local

        mocker.patch("src.config.BATCH_ID", "test_batch_qc")
        mock_to_csv = mocker.patch.object(pd.DataFrame, "to_csv")

        df = pd.DataFrame({"Field": ["RATE"], "Percentage_Filled": [100.0]})
        write_local(df, "", "run_20241114_10-00_test", "qc_qa_stats")

        # Verify to_csv was called
        mock_to_csv.assert_called_once()
        call_args = mock_to_csv.call_args[0][0]

        # Check output path includes qa_qc_output folder and correct filename
        assert "qa_qc_output" in call_args
        assert "QC-QA-STATS.csv" in call_args