feat: integrate PluginStore and MinIO repository into model manager activities
- Added PluginStore integration for model management. - Replaced StorageRepository with MinIORepository in Activities, Cleanup, and Training classes. - Updated training logic to handle validation files and improved data management. - Enhanced configuration for MinIO and PluginStore in connectors. - Removed deprecated model repository and storage repository files. - Updated environment variable handling for new configurations.
This commit is contained in:
vitor-aignosi
@@ -1,982 +0,0 @@
|
||||
"""Unit tests for ModelRepository with 100% coverage."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cleanup_temp_directories():
|
||||
"""Clean up temporary directories after each test."""
|
||||
# Get the temp directory path
|
||||
current_file_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
model_manager_dir = os.path.dirname(os.path.dirname(os.path.dirname(current_file_dir)))
|
||||
temp_dir = os.path.join(model_manager_dir, 'reports', 'temp')
|
||||
|
||||
# Run the test
|
||||
yield
|
||||
|
||||
# Clean up after test
|
||||
if os.path.exists(temp_dir):
|
||||
for item in os.listdir(temp_dir):
|
||||
item_path = os.path.join(temp_dir, item)
|
||||
if os.path.isdir(item_path) and item.startswith('test_run_'):
|
||||
try:
|
||||
shutil.rmtree(item_path)
|
||||
except (OSError, PermissionError):
|
||||
# Ignore cleanup errors
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
"""Create a mock logger."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_train_result():
|
||||
"""Create a mock TrainModelResult."""
|
||||
result = MagicMock()
|
||||
result.params = MagicMock()
|
||||
result.params.experiment_name = 'test_experiment'
|
||||
result.params.experiment_run_id = 1
|
||||
result.params.target_variable = 'target'
|
||||
result.params.variable_columns = ['var1', 'var2']
|
||||
result.params.lag_train = 5
|
||||
result.params.lag_val = 3
|
||||
result.params.window = 10
|
||||
result.params.low_lim = {'var1': 0.0}
|
||||
result.params.upp_lim = {'var1': 10.0}
|
||||
result.params.include_ar = False
|
||||
result.params.train_size = 80
|
||||
result.params.removed_intervals = []
|
||||
result.params.rem_static_win = True
|
||||
result.params.static_threshold = None
|
||||
result.run_name = 'test_run'
|
||||
result.run_dir = '/tmp/test_run' # noqa: S108
|
||||
result.report_path = '/tmp/test_run/report.html' # noqa: S108
|
||||
result.train_data_path = '/tmp/test_run/train_data.csv' # noqa: S108
|
||||
result.test_data_path = '/tmp/test_run/test_data.csv' # noqa: S108
|
||||
result.mse_val = 0.5
|
||||
result.r2_val = 0.9
|
||||
result.mae_val = 0.3
|
||||
result.scaler_dict = {'scaler': 'standard'}
|
||||
result.process_data = MagicMock()
|
||||
result.regr = MagicMock()
|
||||
result.regr.predict = MagicMock(return_value=np.array([1.0, 2.0, 3.0]))
|
||||
result.x_train = pd.DataFrame({'var1': [1, 2, 3]})
|
||||
result.y_train = pd.Series([1.0, 2.0, 3.0], name='target')
|
||||
result.x_test = pd.DataFrame({'var1': [4, 5, 6]})
|
||||
result.y_test = pd.Series([4.0, 5.0, 6.0], name='target')
|
||||
result.y_pred = np.array([4.1, 5.1, 6.1])
|
||||
return result
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_model_repository_init(mock_model_serving_class, mock_logger):
|
||||
"""Test ModelRepository initialization."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_model_serving_class.assert_called_once_with(
|
||||
tracking_uri='http://mlflow.test', username='user', password='pass'
|
||||
)
|
||||
assert repo.model_serving is mock_model_serving_instance
|
||||
assert repo.logger is mock_logger
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_save_model_success(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test save_model successfully saves model."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Reset mock after initialization to focus on method-specific calls
|
||||
mock_logger.reset_mock()
|
||||
|
||||
repo._get_next_run_name = MagicMock(return_value='test_experiment-1')
|
||||
repo._generate_artifacts = MagicMock(return_value=mock_train_result)
|
||||
repo._save_run = MagicMock()
|
||||
|
||||
result = repo.save_model(mock_train_result)
|
||||
|
||||
repo._get_next_run_name.assert_called_once_with('test_experiment')
|
||||
repo._generate_artifacts.assert_called_once()
|
||||
repo._save_run.assert_called_once_with(mock_train_result)
|
||||
mock_logger.info.assert_called_once()
|
||||
assert result is mock_train_result
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.os.path.exists')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.rmtree')
|
||||
def test_cleanup_run_directory_exists(
|
||||
mock_rmtree, mock_exists, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test cleanup_run_directory when directory exists."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Reset mock after initialization to focus on method-specific calls
|
||||
mock_logger.reset_mock()
|
||||
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo.cleanup_run_directory('/tmp/test_run') # noqa: S108
|
||||
|
||||
mock_exists.assert_called_once_with('/tmp/test_run') # noqa: S108
|
||||
mock_rmtree.assert_called_once_with('/tmp/test_run') # noqa: S108
|
||||
mock_logger.info.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.os.path.exists')
|
||||
def test_cleanup_run_directory_not_exists(mock_exists, mock_model_serving_class, mock_logger):
|
||||
"""Test cleanup_run_directory when directory doesn't exist."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Reset mock after initialization to focus on method-specific calls
|
||||
mock_logger.reset_mock()
|
||||
|
||||
mock_exists.return_value = False
|
||||
|
||||
repo.cleanup_run_directory('/tmp/test_run') # noqa: S108
|
||||
|
||||
mock_exists.assert_called_once_with('/tmp/test_run') # noqa: S108
|
||||
mock_logger.info.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_cleanup_run_directory_empty_path(mock_model_serving_class, mock_logger):
|
||||
"""Test cleanup_run_directory with empty path."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Reset mock after initialization to focus on method-specific calls
|
||||
mock_logger.reset_mock()
|
||||
|
||||
repo.cleanup_run_directory('')
|
||||
|
||||
mock_logger.info.assert_called_once_with('No run directory specified, skipping cleanup')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_get_next_run_name_no_existing_runs(mock_model_serving_class, mock_logger):
|
||||
"""Test _get_next_run_name when no runs exist."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_instance.search_runs_by_name.return_value = []
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
result = repo._get_next_run_name('test_experiment')
|
||||
|
||||
assert result == 'test_experiment-1'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_get_next_run_name_with_existing_runs(mock_model_serving_class, mock_logger):
|
||||
"""Test _get_next_run_name when runs exist."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_instance.search_runs_by_name.return_value = [
|
||||
MagicMock(),
|
||||
MagicMock(),
|
||||
MagicMock(),
|
||||
]
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
result = repo._get_next_run_name('test_experiment')
|
||||
|
||||
assert result == 'test_experiment-4'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_get_reports_directory(mock_model_serving_class, mock_logger):
|
||||
"""Test _get_reports_directory returns correct path."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
result = repo._get_reports_directory()
|
||||
|
||||
assert result.endswith(os.path.join('model_manager', 'reports'))
|
||||
assert os.path.isabs(result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_success(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test _init_artifacts_data successfully prepares data."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
reference_data, current_data = repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
# Check reference data
|
||||
assert 'target' in reference_data.columns
|
||||
assert 'prediction' in reference_data.columns
|
||||
assert len(reference_data) == 3
|
||||
|
||||
# Check current data
|
||||
assert 'target' in current_data.columns
|
||||
assert 'prediction' in current_data.columns
|
||||
assert len(current_data) == 3
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_empty_x_train(
|
||||
mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _init_artifacts_data raises ValueError when x_train is empty."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.x_train = pd.DataFrame()
|
||||
|
||||
with pytest.raises(ValueError, match='Training features .* are empty'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_empty_y_train(
|
||||
mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _init_artifacts_data raises ValueError when y_train is empty."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.y_train = pd.Series(dtype=float)
|
||||
|
||||
with pytest.raises(ValueError, match='Training target .* is empty'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_none_y_pred(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test _init_artifacts_data raises ValueError when y_pred is None."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.y_pred = None
|
||||
|
||||
with pytest.raises(ValueError, match='Test predictions .* are None'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_none_y_train_pred(
|
||||
mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _init_artifacts_data raises ValueError when y_train_pred is None."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.y_train_pred = None
|
||||
|
||||
with pytest.raises(ValueError, match='Training predictions .* are None'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.datetime')
|
||||
@patch('model_manager.utils.repository.model_repository.makedirs')
|
||||
def test_create_run_directory_success(
|
||||
mock_makedirs, mock_datetime, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test _create_run_directory creates directory successfully."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_datetime.now.return_value.strftime.return_value = '20240101_120000_123456'
|
||||
|
||||
result = repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108
|
||||
|
||||
expected_path = os.path.normpath(
|
||||
os.path.join('/tmp/reports', 'temp', 'test_run_20240101_120000_123456') # noqa: S108
|
||||
)
|
||||
assert os.path.normpath(result) == expected_path
|
||||
mock_makedirs.assert_called_once()
|
||||
call_path = mock_makedirs.call_args[0][0]
|
||||
assert os.path.normpath(call_path) == expected_path
|
||||
assert mock_makedirs.call_args[1] == {'exist_ok': True}
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.makedirs')
|
||||
def test_create_run_directory_permission_error(
|
||||
mock_makedirs, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test _create_run_directory raises PermissionError."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_makedirs.side_effect = PermissionError('Permission denied')
|
||||
|
||||
with pytest.raises(PermissionError, match='Permission denied when creating directory'):
|
||||
repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.copy')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_setup_run_directory_success(mock_open, mock_copy, mock_model_serving_class, mock_logger):
|
||||
"""Test _setup_run_directory creates files successfully."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
|
||||
|
||||
# Check that empty files were created
|
||||
assert mock_open.call_count == 3
|
||||
mock_copy.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.json.dump')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_with_equation(
|
||||
mock_open,
|
||||
mock_reports_class,
|
||||
mock_json_dump,
|
||||
mock_model_serving_class,
|
||||
mock_logger,
|
||||
mock_train_result,
|
||||
):
|
||||
"""Test _generate_report creates equation JSON artifact."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Add equation to train result
|
||||
mock_train_result.equation = {
|
||||
'target_variable': 'target',
|
||||
'coefficients': {'var1': 1.5, 'var2': -0.75},
|
||||
'intercept': 10.5,
|
||||
'equation_string': 'target = 10.5 + 1.5 * var1 + -0.75 * var2',
|
||||
'latex_equation': 'target = 10.5 + 1.5 \\cdot var1 + -0.75 \\cdot var2',
|
||||
'model_type': 'Linear Regression',
|
||||
}
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
# Mock DataFrame.to_csv to avoid file I/O
|
||||
with patch.object(pd.DataFrame, 'to_csv'):
|
||||
result = repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
# Verify equation path was set
|
||||
assert result.equation_path == os.path.join(mock_train_result.run_dir, 'model_equation.json')
|
||||
|
||||
# Verify JSON was written
|
||||
mock_json_dump.assert_called()
|
||||
call_args = mock_json_dump.call_args
|
||||
assert call_args[0][0] == mock_train_result.equation
|
||||
assert call_args[1]['indent'] == 2
|
||||
assert call_args[1]['ensure_ascii'] is False
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_without_equation(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report works without equation."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# No equation
|
||||
mock_train_result.equation = None
|
||||
# Remove equation_path if it exists from fixture
|
||||
if hasattr(mock_train_result, 'equation_path'):
|
||||
del mock_train_result.equation_path
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
# Mock DataFrame.to_csv to avoid file I/O
|
||||
with patch.object(pd.DataFrame, 'to_csv'):
|
||||
result = repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
# Verify equation section was not executed (equation_path not set)
|
||||
# Since equation is None, the equation block should not run
|
||||
assert result == mock_train_result
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_with_equation(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run logs equation artifact."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Set equation path
|
||||
mock_train_result.equation_path = '/tmp/test_run/model_equation.json' # noqa: S108
|
||||
|
||||
# Mock all path.exists calls to return True
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
# Verify equation artifact was logged
|
||||
logged_artifacts = [
|
||||
call[0][0] for call in mock_model_serving_instance.log_artifact.call_args_list
|
||||
]
|
||||
assert '/tmp/test_run/model_equation.json' in logged_artifacts # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_with_static_threshold_value(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run logs static_threshold when rem_static_win is True and value is set."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Set static_threshold to a specific value
|
||||
mock_train_result.params.rem_static_win = True
|
||||
mock_train_result.params.static_threshold = 500
|
||||
mock_train_result.equation_path = None
|
||||
|
||||
# Mock all path.exists calls to return True
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
# Verify static_threshold was logged with the correct value
|
||||
log_param_calls = {
|
||||
call[0][0]: call[0][1] for call in mock_model_serving_instance.log_param.call_args_list
|
||||
}
|
||||
assert log_param_calls['static_threshold'] == 500
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_with_rem_static_win_false(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run logs static_threshold as None when rem_static_win is False."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Set rem_static_win to False
|
||||
mock_train_result.params.rem_static_win = False
|
||||
mock_train_result.params.static_threshold = 500 # Should be ignored
|
||||
mock_train_result.equation_path = None
|
||||
|
||||
# Mock all path.exists calls to return True
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
# Verify static_threshold was logged as None
|
||||
log_param_calls = {
|
||||
call[0][0]: call[0][1] for call in mock_model_serving_instance.log_param.call_args_list
|
||||
}
|
||||
assert log_param_calls['static_threshold'] is None
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_without_equation(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run works without equation."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# No equation
|
||||
mock_train_result.equation_path = None
|
||||
|
||||
# Mock path.exists to return True for required artifacts
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
# Verify only 3 artifacts were logged (report, train_data, test_data)
|
||||
assert mock_model_serving_instance.log_artifact.call_count == 3
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_missing_report(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run raises ValueError when report is missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock report doesn't exist
|
||||
def exists_side_effect(path):
|
||||
return not path.endswith('report.html')
|
||||
|
||||
mock_exists.side_effect = exists_side_effect
|
||||
|
||||
with pytest.raises(ValueError, match='Report file does not exist'):
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_none_metrics(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run raises ValueError when metrics are None."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.mse_val = None
|
||||
# Mock all paths exist so we reach the metrics check
|
||||
mock_exists.return_value = True
|
||||
|
||||
with pytest.raises(ValueError, match='One or more metrics .* are None'):
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_missing_train_data(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run raises ValueError when train data is missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock train_data doesn't exist
|
||||
def exists_side_effect(path):
|
||||
return not path.endswith('train_data.csv')
|
||||
|
||||
mock_exists.side_effect = exists_side_effect
|
||||
|
||||
with pytest.raises(ValueError, match='Training data file does not exist'):
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_missing_test_data(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run raises ValueError when test data is missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock test_data doesn't exist
|
||||
def exists_side_effect(path):
|
||||
if path.endswith('test_data.csv'):
|
||||
return False
|
||||
return True
|
||||
|
||||
mock_exists.side_effect = exists_side_effect
|
||||
|
||||
with pytest.raises(ValueError, match='Test data file does not exist'):
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_empty_x_test(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test _init_artifacts_data raises ValueError when x_test is empty."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.x_test = pd.DataFrame()
|
||||
|
||||
with pytest.raises(ValueError, match='Test features .* are empty'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_empty_y_test(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test _init_artifacts_data raises ValueError when y_test is empty."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.y_test = pd.Series(dtype=float)
|
||||
|
||||
with pytest.raises(ValueError, match='Test target .* is empty'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.makedirs')
|
||||
def test_create_run_directory_os_error(mock_makedirs, mock_model_serving_class, mock_logger):
|
||||
"""Test _create_run_directory raises OSError."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_makedirs.side_effect = OSError('Disk full')
|
||||
|
||||
with pytest.raises(OSError, match='Failed to create directory'):
|
||||
repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.copy')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_setup_run_directory_file_not_found(
|
||||
mock_open, mock_copy, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test _setup_run_directory raises FileNotFoundError when header missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_copy.side_effect = FileNotFoundError('Header not found')
|
||||
|
||||
with pytest.raises(FileNotFoundError, match='Header file not found'):
|
||||
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.copy')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_setup_run_directory_permission_error(
|
||||
mock_open, mock_copy, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test _setup_run_directory raises PermissionError."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_open.side_effect = PermissionError('Permission denied')
|
||||
|
||||
with pytest.raises(PermissionError, match='Permission denied when setting up directory'):
|
||||
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.copy')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_setup_run_directory_os_error(mock_open, mock_copy, mock_model_serving_class, mock_logger):
|
||||
"""Test _setup_run_directory raises OSError."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_open.side_effect = OSError('Disk error')
|
||||
|
||||
with pytest.raises(OSError, match='Failed to setup run directory'):
|
||||
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_value_error(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report raises ValueError on invalid data."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Create data that can't be converted to float64
|
||||
reference_data = pd.DataFrame({'var1': ['invalid', 'data']})
|
||||
current_data = pd.DataFrame({'var1': [1, 2]})
|
||||
|
||||
with pytest.raises(ValueError, match='Failed to convert data to float64'):
|
||||
repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_permission_error(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report raises PermissionError on write failure."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
# Mock Reports to raise PermissionError
|
||||
mock_reports_class.side_effect = PermissionError('Permission denied')
|
||||
|
||||
with pytest.raises(PermissionError, match='Permission denied when writing report files'):
|
||||
repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_os_error(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report raises OSError on write failure."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
# Mock Reports to raise OSError
|
||||
mock_reports_class.side_effect = OSError('Disk error')
|
||||
|
||||
with pytest.raises(OSError, match='Failed to generate report'):
|
||||
repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_none_run_dir(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report raises ValueError when run_dir is None."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.run_dir = None
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
with pytest.raises(ValueError, match='run_dir is not set'):
|
||||
repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_generate_artifacts_no_run_name(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_artifacts raises ValueError when run_name is not set."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.run_name = None
|
||||
mock_exists.return_value = True # Mock reports directory exists
|
||||
|
||||
# Mock _create_run_directory to avoid creating real directories
|
||||
with patch.object(repo, '_create_run_directory') as mock_create_dir:
|
||||
mock_create_dir.return_value = '/mock/run/dir'
|
||||
|
||||
with pytest.raises(ValueError, match='run_name must be set'):
|
||||
repo._generate_artifacts(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_generate_artifacts_reports_dir_not_found(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_artifacts raises FileNotFoundError when reports dir missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock reports directory doesn't exist
|
||||
mock_exists.return_value = False
|
||||
|
||||
with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
|
||||
repo._generate_artifacts(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_generate_artifacts_header_not_found(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_artifacts raises FileNotFoundError when header missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock: reports dir exists, but header doesn't
|
||||
def exists_side_effect(path):
|
||||
if path.endswith('header.html'):
|
||||
return False
|
||||
return True
|
||||
|
||||
mock_exists.side_effect = exists_side_effect
|
||||
|
||||
# Mock _create_run_directory to avoid creating real directories
|
||||
with patch.object(repo, '_create_run_directory') as mock_create_dir:
|
||||
mock_create_dir.return_value = '/mock/run/dir'
|
||||
|
||||
with pytest.raises(FileNotFoundError, match='Header file does not exist'):
|
||||
repo._generate_artifacts(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
@patch('model_manager.utils.repository.model_repository.path.join')
|
||||
def test_generate_artifacts_success(
|
||||
mock_join, mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_artifacts success case covering lines 147-148."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock path.join to return predictable paths
|
||||
def join_side_effect(*args):
|
||||
return '/'.join(args)
|
||||
|
||||
mock_join.side_effect = join_side_effect
|
||||
mock_exists.return_value = True # Both reports dir and header.html exist
|
||||
|
||||
# Mock the internal methods to avoid actual file operations
|
||||
with (
|
||||
patch.object(repo, '_setup_run_directory') as mock_setup,
|
||||
patch.object(repo, '_generate_report') as mock_generate_report,
|
||||
patch.object(repo, '_init_artifacts_data') as mock_init_data,
|
||||
patch.object(repo, '_get_reports_directory') as mock_get_reports_dir,
|
||||
patch.object(repo, '_create_run_directory') as mock_create_run_dir,
|
||||
):
|
||||
# Setup mocks
|
||||
mock_init_data.return_value = (pd.DataFrame(), pd.DataFrame())
|
||||
mock_get_reports_dir.return_value = '/reports'
|
||||
mock_create_run_dir.return_value = '/reports/run_1'
|
||||
mock_generate_report.return_value = mock_train_result
|
||||
|
||||
# Call the method
|
||||
result = repo._generate_artifacts(mock_train_result)
|
||||
|
||||
# Verify the methods on lines 147-148 were called
|
||||
mock_setup.assert_called_once_with('/reports/run_1', '/reports/header.html')
|
||||
mock_generate_report.assert_called_once()
|
||||
|
||||
# Verify result
|
||||
assert result == mock_train_result
|
||||
@@ -1,513 +0,0 @@
|
||||
"""Unit tests for StorageRepository class."""
|
||||
|
||||
from io import BytesIO
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from botocore.exceptions import ClientError
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
"""Create a mock logger for testing."""
|
||||
logger = Mock()
|
||||
logger.info = Mock()
|
||||
logger.error = Mock()
|
||||
logger.warning = Mock()
|
||||
return logger
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage_config():
|
||||
"""Create storage repository configuration."""
|
||||
return {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'test_access_key',
|
||||
'secret_key': 'test_secret_key',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'standard',
|
||||
'connect_timeout': 30,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_initialization(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository initialization with correct parameters."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Verify attributes are set correctly
|
||||
assert repo.endpoint_url == storage_config['endpoint_url']
|
||||
assert repo.access_key == storage_config['access_key']
|
||||
assert repo.secret_key == storage_config['secret_key']
|
||||
assert repo.region == storage_config['region']
|
||||
assert repo.use_ssl == storage_config['use_ssl']
|
||||
assert repo.max_retry_attempts == storage_config['max_retry_attempts']
|
||||
assert repo.retry_mode == storage_config['retry_mode']
|
||||
assert repo.connect_timeout == storage_config['connect_timeout']
|
||||
assert repo.read_timeout == storage_config['read_timeout']
|
||||
assert repo.logger == mock_logger
|
||||
|
||||
# Verify boto3 client was created
|
||||
mock_boto3.client.assert_called_once()
|
||||
call_args = mock_boto3.client.call_args
|
||||
|
||||
assert call_args[0][0] == 's3'
|
||||
assert call_args[1]['endpoint_url'] == storage_config['endpoint_url']
|
||||
assert call_args[1]['aws_access_key_id'] == storage_config['access_key']
|
||||
assert call_args[1]['aws_secret_access_key'] == storage_config['secret_key']
|
||||
assert call_args[1]['use_ssl'] == storage_config['use_ssl']
|
||||
|
||||
# Verify logger was called
|
||||
mock_logger.info.assert_called()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_boto_config(mock_boto3, mock_logger, storage_config):
|
||||
"""Test that boto3 Config is created with correct retry settings."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Verify Config was passed with correct settings
|
||||
call_args = mock_boto3.client.call_args
|
||||
boto_config = call_args[1]['config']
|
||||
|
||||
assert boto_config.region_name == storage_config['region']
|
||||
assert boto_config.connect_timeout == storage_config['connect_timeout']
|
||||
assert boto_config.read_timeout == storage_config['read_timeout']
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_success(mock_boto3, mock_logger, storage_config):
|
||||
"""Test successful file fetch from MinIO."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock response
|
||||
file_content = b'test file content'
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = file_content
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Fetch file
|
||||
result = repo.fetch_file('test-bucket', 'test-file.csv')
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, BytesIO)
|
||||
assert result.getvalue() == file_content
|
||||
|
||||
# Verify get_object was called correctly
|
||||
mock_s3_client.get_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.csv')
|
||||
|
||||
# Verify logging
|
||||
assert mock_logger.info.call_count >= 2 # Init + fetch
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_with_large_content(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file with large content."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock response with large content
|
||||
large_content = b'x' * 1024 * 1024 # 1MB
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = large_content
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Fetch file
|
||||
result = repo.fetch_file('test-bucket', 'large-file.bin')
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, BytesIO)
|
||||
assert len(result.getvalue()) == 1024 * 1024
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_empty_content(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file with empty content."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock response with empty content
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = b''
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Fetch file
|
||||
result = repo.fetch_file('test-bucket', 'empty-file.txt')
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, BytesIO)
|
||||
assert result.getvalue() == b''
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_not_found(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file when object doesn't exist."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise NoSuchKey error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.side_effect = ClientError(
|
||||
{'Error': {'Code': 'NoSuchKey', 'Message': 'The specified key does not exist.'}},
|
||||
'GetObject',
|
||||
)
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to fetch non-existent file
|
||||
with pytest.raises(ClientError) as exc_info:
|
||||
repo.fetch_file('test-bucket', 'non-existent.csv')
|
||||
|
||||
assert exc_info.value.response['Error']['Code'] == 'NoSuchKey'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_access_denied(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file when access is denied."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise AccessDenied error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.side_effect = ClientError(
|
||||
{'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}, 'GetObject'
|
||||
)
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to fetch file without permissions
|
||||
with pytest.raises(ClientError) as exc_info:
|
||||
repo.fetch_file('test-bucket', 'protected-file.csv')
|
||||
|
||||
assert exc_info.value.response['Error']['Code'] == 'AccessDenied'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_network_error(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file when network error occurs."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise network error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.side_effect = ConnectionError('Network unreachable')
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to fetch file with network error
|
||||
with pytest.raises(ConnectionError):
|
||||
repo.fetch_file('test-bucket', 'test-file.csv')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_success(mock_boto3, mock_logger, storage_config):
|
||||
"""Test successful file deletion from MinIO."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete file
|
||||
repo.delete_file('test-bucket', 'test-file.csv')
|
||||
|
||||
# Verify delete_object was called correctly
|
||||
mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.csv')
|
||||
|
||||
# Verify logging
|
||||
assert mock_logger.info.call_count >= 2 # Init + delete
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_non_existent(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file that doesn't exist (should succeed silently in S3)."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# S3/MinIO delete is idempotent - deleting non-existent file succeeds
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete non-existent file (should succeed)
|
||||
repo.delete_file('test-bucket', 'non-existent.csv')
|
||||
|
||||
mock_s3_client.delete_object.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_access_denied(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file when access is denied."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise AccessDenied error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.side_effect = ClientError(
|
||||
{'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}, 'DeleteObject'
|
||||
)
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to delete file without permissions
|
||||
with pytest.raises(ClientError) as exc_info:
|
||||
repo.delete_file('test-bucket', 'protected-file.csv')
|
||||
|
||||
assert exc_info.value.response['Error']['Code'] == 'AccessDenied'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_network_error(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file when network error occurs."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise network error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.side_effect = ConnectionError('Network unreachable')
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to delete file with network error
|
||||
with pytest.raises(ConnectionError):
|
||||
repo.delete_file('test-bucket', 'test-file.csv')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_with_ssl(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository initialization with SSL enabled."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
storage_config['use_ssl'] = True
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
assert repo.use_ssl is True
|
||||
|
||||
# Verify boto3 client was created with use_ssl=True
|
||||
call_args = mock_boto3.client.call_args
|
||||
assert call_args[1]['use_ssl'] is True
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_custom_timeouts(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with custom timeout values."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
storage_config['connect_timeout'] = 10
|
||||
storage_config['read_timeout'] = 120
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
assert repo.connect_timeout == 10
|
||||
assert repo.read_timeout == 120
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_custom_retry_mode(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with different retry modes."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
for retry_mode in ['standard', 'legacy', 'adaptive']:
|
||||
storage_config['retry_mode'] = retry_mode
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
assert repo.retry_mode == retry_mode
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_custom_max_retries(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with different max retry attempts."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
storage_config['max_retry_attempts'] = 5
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
assert repo.max_retry_attempts == 5
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_with_special_characters(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file with special characters in name."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
file_content = b'test content'
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = file_content
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Fetch file with special characters
|
||||
special_filename = 'test file (2023-01-01) #1.csv'
|
||||
result = repo.fetch_file('test-bucket', special_filename)
|
||||
|
||||
assert isinstance(result, BytesIO)
|
||||
mock_s3_client.get_object.assert_called_once_with(Bucket='test-bucket', Key=special_filename)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_with_path_separators(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file with path separators in object key."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete file with path separators
|
||||
file_path = 'data/2023/01/test-file.csv'
|
||||
repo.delete_file('test-bucket', file_path)
|
||||
|
||||
mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key=file_path)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_different_regions(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with different AWS regions."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
regions = ['us-west-1', 'eu-central-1', 'ap-southeast-1']
|
||||
|
||||
for region in regions:
|
||||
storage_config['region'] = region
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
assert repo.region == region
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_logs_file_size(mock_boto3, mock_logger, storage_config):
|
||||
"""Test that fetch_file logs the file size."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
file_content = b'x' * 12345
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = file_content
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
repo.fetch_file('test-bucket', 'test-file.csv')
|
||||
|
||||
# Verify logging includes file size
|
||||
log_calls = [str(call) for call in mock_logger.info.call_args_list]
|
||||
assert any('12345 bytes' in str(call) for call in log_calls)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_close_method(mock_boto3, mock_logger, storage_config):
|
||||
"""Test that the close method calls the underlying client's close method."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
repo.close()
|
||||
|
||||
mock_s3_client.close.assert_called_once()
|
||||
mock_logger.info.assert_called_with('MinIO client closed')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_list_bucket_objects_with_pagination(mock_boto3, mock_logger, storage_config):
|
||||
"""Test list_bucket_objects with a paginated response."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_paginator = Mock()
|
||||
page1 = {
|
||||
'Contents': [
|
||||
{'Key': 'file1.txt'},
|
||||
{'Key': 'file2.txt'},
|
||||
]
|
||||
}
|
||||
page2 = {
|
||||
'Contents': [
|
||||
{'Key': 'file3.txt'},
|
||||
]
|
||||
}
|
||||
page3 = {}
|
||||
|
||||
mock_paginator.paginate.return_value = [page1, page2, page3]
|
||||
mock_s3_client.get_paginator.return_value = mock_paginator
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
objects = repo.list_bucket_objects('test-bucket', max_keys=2)
|
||||
|
||||
assert objects == ['file1.txt', 'file2.txt', 'file3.txt']
|
||||
assert len(objects) == 3
|
||||
mock_s3_client.get_paginator.assert_called_once_with('list_objects_v2')
|
||||
mock_paginator.paginate.assert_called_once_with(Bucket='test-bucket', MaxKeys=2)
|
||||
mock_logger.info.assert_any_call('Listed 3 objects from bucket test-bucket')
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user