Siddhant Medar
ab71bd8dae
Task/unit tests investment and client * move io_utils.py * add __init__.py * refactor: move string_funcs and update import paths to use utils namespace * rename string_funcs to string_utils * move and rename table_utils.py * refactor: rename claude_funcs to llm_utils in multiple files * update unit tests * added unit tests for src/string_utils.py * added unit tests for src/table_utils.py * added unit tests for src/llm_utils.py * added unit tests for src/io_utils.py * Merge branch 'main' into task/unit-tests-investment-and-client; TODO: fix unit tests * Unit tests to fix * fix unit tests * fix invalid references Approved-by: Alex Galarce
265 lines
11 KiB
Python
265 lines
11 KiB
Python
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() |