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