Alex Galarce
0562b90646
io_utils documentation * add docstring to read_local function and handle file read errors gracefully by returning None on failure * update read_local function to support .xlsx files and improve error handling * add docstring to read_s3 function to clarify its purpose and return value * refactor read_s3 and read_s3_csv functions to improve error handling and add type hints * improve error handling in read_s3_csv function for CSV file parsing * enhance read_input function with error handling and detailed docstring * Add TODOs for read_input_csv and filter_already_processed functions; add type hints and docstring for remove_txt_extension * enhance write_local function documentation; make docstrings more consistent when there is a "notes" section * enhance write_local and write_s3 functions; improve documentation, add error output type handling, and ensure consistent CSV writing options * refactor test_io_utils.py; reorganize imports and enhance read_input_local test with additional mock for file existence Approved-by: Katon Minhas
312 lines
12 KiB
Python
312 lines
12 KiB
Python
from io import StringIO
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
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,
|
|
)
|
|
|
|
|
|
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("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()
|
|
# 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()
|