284 lines
9.6 KiB
Python
284 lines
9.6 KiB
Python
"""Unit tests for Training activity."""
|
|
|
|
from io import BytesIO
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from pytest import mark
|
|
|
|
from model_manager.activities.training import Training
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.training.TrainingRepository')
|
|
async def test_train_model_success(mock_training_repository_class):
|
|
"""Test successful model training."""
|
|
# Create mock repository instance
|
|
mock_repository = MagicMock()
|
|
mock_training_repository_class.return_value = mock_repository
|
|
|
|
# Create mock train result
|
|
mock_train_result = MagicMock(spec=TrainModelResult)
|
|
mock_train_result.mse_val = 0.5
|
|
mock_train_result.mae_val = 0.3
|
|
mock_train_result.r2_val = 0.95
|
|
|
|
mock_final_result = MagicMock(spec=TrainModelResult)
|
|
mock_final_result.mse_val = 0.5
|
|
mock_final_result.mae_val = 0.3
|
|
mock_final_result.r2_val = 0.95
|
|
|
|
# Setup repository mocks
|
|
mock_repository.train.return_value = mock_train_result
|
|
mock_repository.after_train_calculation.return_value = mock_final_result
|
|
|
|
# Create Training instance
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
# Mock inherited methods
|
|
training.info = MagicMock()
|
|
|
|
# Test data
|
|
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
|
|
train_params_dict = {
|
|
'experiment_run_id': 123,
|
|
'target_variable': 'price',
|
|
'variable_columns': ['feature1', 'feature2'],
|
|
'train_size': 80,
|
|
'shuffle': True,
|
|
'use_scaler': True,
|
|
'include_ar': False,
|
|
'bucket_name': 'test-bucket',
|
|
'file_name': 'test.csv',
|
|
'line_separator': '\n',
|
|
'decimal_separator': '.',
|
|
'lag_train': 1,
|
|
'lag_val': 1,
|
|
'rem_static_win': False,
|
|
'low_lim': {'feature1': 0.0, 'feature2': 0.0},
|
|
'upp_lim': {'feature1': 100.0, 'feature2': 100.0},
|
|
'window': 10,
|
|
'experiment_name': 'test_experiment',
|
|
'removed_intervals': [],
|
|
}
|
|
|
|
input_data = {
|
|
'metadata': {'workflow_id': 'test-123'},
|
|
'uploaded_file': uploaded_file,
|
|
'train_params': train_params_dict,
|
|
}
|
|
|
|
# Execute
|
|
result = await training.train_model(input_data)
|
|
|
|
# Assertions
|
|
assert result['success'] is True
|
|
assert result['result'] == mock_final_result
|
|
assert result['error_message'] is None
|
|
|
|
# Verify repository calls
|
|
mock_repository.train.assert_called_once()
|
|
mock_repository.after_train_calculation.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.training.TrainingRepository')
|
|
async def test_train_model_invalid_file_type(mock_training_repository_class):
|
|
"""Test training with invalid file type."""
|
|
mock_repository = MagicMock()
|
|
mock_training_repository_class.return_value = mock_repository
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
# Invalid file type (string instead of BytesIO)
|
|
input_data = {
|
|
'metadata': {},
|
|
'uploaded_file': 'not_a_bytesio',
|
|
'train_params': {
|
|
'experiment_run_id': 123,
|
|
'target_variable': 'price',
|
|
'variable_columns': ['feature1'],
|
|
'train_size': 80,
|
|
'shuffle': True,
|
|
'use_scaler': False,
|
|
'include_ar': False,
|
|
'bucket_name': 'test',
|
|
'file_name': 'test.csv',
|
|
'line_separator': '\n',
|
|
'decimal_separator': '.',
|
|
'lag_train': 1,
|
|
'lag_val': 1,
|
|
'rem_static_win': False,
|
|
'low_lim': {'feature1': 0.0},
|
|
'upp_lim': {'feature1': 100.0},
|
|
'window': 10,
|
|
'experiment_name': 'test_experiment',
|
|
'removed_intervals': [],
|
|
},
|
|
}
|
|
|
|
result = await training.train_model(input_data)
|
|
|
|
assert result['success'] is False
|
|
assert result['result'] is None
|
|
assert 'uploaded_file must be BytesIO' in result['error_message']
|
|
# Verify notification was sent (via BaseActivity)
|
|
notification_handler.send_notification.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.training.TrainingRepository')
|
|
async def test_train_model_training_error(mock_training_repository_class):
|
|
"""Test training failure during model training."""
|
|
mock_repository = MagicMock()
|
|
mock_training_repository_class.return_value = mock_repository
|
|
|
|
# Setup repository to raise error
|
|
mock_repository.train.side_effect = ValueError('Training data is empty')
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
uploaded_file = BytesIO(b'test,data\n')
|
|
input_data = {
|
|
'metadata': {'workflow_id': 'test-456'},
|
|
'uploaded_file': uploaded_file,
|
|
'train_params': {
|
|
'experiment_run_id': 456,
|
|
'target_variable': 'price',
|
|
'variable_columns': ['feature1'],
|
|
'train_size': 80,
|
|
'shuffle': True,
|
|
'use_scaler': False,
|
|
'include_ar': False,
|
|
'bucket_name': 'test',
|
|
'file_name': 'test.csv',
|
|
'line_separator': '\n',
|
|
'decimal_separator': '.',
|
|
'lag_train': 1,
|
|
'lag_val': 1,
|
|
'rem_static_win': False,
|
|
'low_lim': {'feature1': 0.0},
|
|
'upp_lim': {'feature1': 100.0},
|
|
'window': 10,
|
|
'experiment_name': 'test_experiment',
|
|
'removed_intervals': [],
|
|
},
|
|
}
|
|
|
|
result = await training.train_model(input_data)
|
|
|
|
assert result['success'] is False
|
|
assert result['result'] is None
|
|
assert 'Training data is empty' in result['error_message']
|
|
# Verify notification was sent (via BaseActivity)
|
|
notification_handler.send_notification.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.training.TrainingRepository')
|
|
async def test_train_model_sends_notification_on_error(mock_training_repository_class):
|
|
"""Test that notification is sent when training fails."""
|
|
mock_repository = MagicMock()
|
|
mock_training_repository_class.return_value = mock_repository
|
|
|
|
mock_repository.train.side_effect = Exception('Database connection failed')
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
uploaded_file = BytesIO(b'test,data\n1,2')
|
|
input_data = {
|
|
'metadata': {'workflow_id': 'test-789', 'experiment_run_id': 789},
|
|
'uploaded_file': uploaded_file,
|
|
'train_params': {
|
|
'experiment_run_id': 789,
|
|
'target_variable': 'price',
|
|
'variable_columns': ['feature1'],
|
|
'train_size': 80,
|
|
'shuffle': True,
|
|
'use_scaler': False,
|
|
'include_ar': False,
|
|
'bucket_name': 'test',
|
|
'file_name': 'test.csv',
|
|
'line_separator': '\n',
|
|
'decimal_separator': '.',
|
|
'lag_train': 1,
|
|
'lag_val': 1,
|
|
'rem_static_win': False,
|
|
'low_lim': {'feature1': 0.0},
|
|
'upp_lim': {'feature1': 100.0},
|
|
'window': 10,
|
|
'experiment_name': 'test_experiment',
|
|
'removed_intervals': [],
|
|
},
|
|
}
|
|
|
|
result = await training.train_model(input_data)
|
|
|
|
# Verify notification was sent (via BaseActivity)
|
|
notification_handler.send_notification.assert_called_once()
|
|
|
|
# Verify result - the important part is that error was caught and returned
|
|
assert result['success'] is False
|
|
assert result['result'] is None
|
|
assert 'Database connection failed' in result['error_message']
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.training.TrainingRepository')
|
|
async def test_train_model_after_calculation_error(mock_training_repository_class):
|
|
"""Test training failure during post-training calculations."""
|
|
mock_repository = MagicMock()
|
|
mock_training_repository_class.return_value = mock_repository
|
|
|
|
# Train succeeds but after_calculation fails
|
|
mock_train_result = MagicMock(spec=TrainModelResult)
|
|
mock_repository.train.return_value = mock_train_result
|
|
mock_repository.after_train_calculation.side_effect = Exception('Metric calculation failed')
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
|
|
input_data = {
|
|
'metadata': {},
|
|
'uploaded_file': uploaded_file,
|
|
'train_params': {
|
|
'experiment_run_id': 999,
|
|
'target_variable': 'price',
|
|
'variable_columns': ['feature1'],
|
|
'train_size': 80,
|
|
'shuffle': True,
|
|
'use_scaler': False,
|
|
'include_ar': False,
|
|
'bucket_name': 'test',
|
|
'file_name': 'test.csv',
|
|
'line_separator': '\n',
|
|
'decimal_separator': '.',
|
|
'lag_train': 1,
|
|
'lag_val': 1,
|
|
'rem_static_win': False,
|
|
'low_lim': {'feature1': 0.0},
|
|
'upp_lim': {'feature1': 100.0},
|
|
'window': 10,
|
|
'experiment_name': 'test_experiment',
|
|
'removed_intervals': [],
|
|
},
|
|
}
|
|
|
|
result = await training.train_model(input_data)
|
|
|
|
assert result['success'] is False
|
|
assert result['result'] is None
|
|
assert 'Metric calculation failed' in result['error_message']
|
|
# Verify notification was sent (via BaseActivity)
|
|
notification_handler.send_notification.assert_called_once()
|