import pytest from io import StringIO import pandas as pd from src.utils.io_utils import ( read_local, list_s3_files, read_s3, read_s3_csv, read_input, read_input_csv, filter_already_processed, remove_txt_extension, read_xlsb, ) class TestIOUtils: # 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("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()