gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1253: Enforce TrainModelParams object in Training activity and tests, removing dict conversion.

Browse Source
This commit is contained in:
Bruno Domingues
2025-10-14 10:03:34 -03:00
parent 014e1cd910
commit 61267ec49d
2 changed files with 167 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

255
tests/activities/test_training.py
View File

@@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch
from pytest import mark
from model_manager.activities.training import Training
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
@@ -42,32 +43,32 @@ async def test_train_model_success(mock_training_repository_class):
# 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': [],
}
train_params = TrainModelParams(
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,
'train_params': train_params,
}
# Execute
@@ -95,30 +96,32 @@ async def test_train_model_invalid_file_type(mock_training_repository_class):
training = Training(logger=logger, notification_handler=notification_handler)
# Invalid file type (string instead of BytesIO)
train_params = TrainModelParams(
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=[],
)
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': [],
},
'train_params': train_params,
}
result = await training.train_model(input_data)
@@ -145,30 +148,32 @@ async def test_train_model_training_error(mock_training_repository_class):
training = Training(logger=logger, notification_handler=notification_handler)
uploaded_file = BytesIO(b'test,data\n')
train_params = TrainModelParams(
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=[],
)
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': [],
},
'train_params': train_params,
}
result = await training.train_model(input_data)
@@ -194,30 +199,32 @@ async def test_train_model_sends_notification_on_error(mock_training_repository_
training = Training(logger=logger, notification_handler=notification_handler)
uploaded_file = BytesIO(b'test,data\n1,2')
train_params = TrainModelParams(
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=[],
)
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': [],
},
'train_params': train_params,
}
result = await training.train_model(input_data)
@@ -248,30 +255,32 @@ async def test_train_model_after_calculation_error(mock_training_repository_clas
training = Training(logger=logger, notification_handler=notification_handler)
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
train_params = TrainModelParams(
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=[],
)
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': [],
},
'train_params': train_params,
}
result = await training.train_model(input_data)
@@ -281,3 +290,37 @@ async def test_train_model_after_calculation_error(mock_training_repository_clas
assert 'Metric calculation failed' 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_invalid_train_params_type(mock_training_repository_class):
"""Test training with invalid train_params type (dict instead of TrainModelParams)."""
mock_repository = MagicMock()
mock_training_repository_class.return_value = mock_repository
logger = MagicMock()
notification_handler = MagicMock()
training = Training(logger=logger, notification_handler=notification_handler)
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
# Invalid train_params type (dict instead of TrainModelParams object)
input_data = {
'metadata': {'workflow_id': 'test-invalid'},
'uploaded_file': uploaded_file,
'train_params': {
'experiment_run_id': 123,
'target_variable': 'price',
'variable_columns': ['feature1'],
}, # This is a dict, not TrainModelParams
}
result = await training.train_model(input_data)
assert result['success'] is False
assert result['result'] is None
assert 'train_params must be TrainModelParams' in result['error_message']
assert 'dict' in result['error_message']
# Verify notification was sent (via BaseActivity)
notification_handler.send_notification.assert_called_once()