342a02d6f7
- Introduced a new activity to load model metadata from the model store. - Refactored training logic to utilize new model metadata and improved parameter handling. - Updated the `TrainModelParams` class to include additional fields for model configuration. - Replaced deprecated utility functions with a custom train-test split implementation. - Removed unused utility functions and cleaned up the data manager repository. - Adjusted experiment tracking to include model-specific metadata in notifications.
252 lines
7.4 KiB
Python
252 lines
7.4 KiB
Python
"""Unit tests for TrainModelResult dataclass."""
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_params():
|
|
"""Create sample TrainModelParams for testing."""
|
|
return TrainModelParams(
|
|
variable_columns=['var1', 'var2'],
|
|
lag_train={'var1': 5, 'var2': 5},
|
|
lag_val={'var1': 3, 'var2': 3},
|
|
target_variable='target',
|
|
rem_static_win=True,
|
|
low_lim={'var1': 0.0, 'var2': 0.0},
|
|
upp_lim={'var1': 100.0, 'var2': 100.0},
|
|
window=10,
|
|
use_scaler=True,
|
|
include_ar=False,
|
|
bucket_name='test-bucket',
|
|
file_name='test-file.csv',
|
|
line_separator='\n',
|
|
decimal_separator='.',
|
|
train_size=80,
|
|
shuffle=True,
|
|
experiment_run_id=123,
|
|
removed_intervals=[],
|
|
model_name='Linear Regression',
|
|
degree=1,
|
|
interaction_only=False,
|
|
nan_treatment='drop',
|
|
start_date=None,
|
|
end_date=None,
|
|
scaler_name='Standard Scaler',
|
|
support_filters={},
|
|
static_threshold=None,
|
|
date_column=None,
|
|
date_format=None,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_dataframes():
|
|
"""Create sample DataFrames for testing."""
|
|
X_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [4, 5, 6]})
|
|
X_test = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10]})
|
|
y_train = pd.DataFrame({'target': [10, 20, 30]})
|
|
y_test = pd.DataFrame({'target': [40, 50]})
|
|
return X_train, X_test, y_train, y_test
|
|
|
|
|
|
def test_train_model_result_creation(sample_params, sample_dataframes):
|
|
"""Test creating TrainModelResult with required fields."""
|
|
x_train, x_test, y_train, y_test = sample_dataframes
|
|
process_data = MagicMock()
|
|
regr = MagicMock()
|
|
|
|
result = TrainModelResult(
|
|
params=sample_params,
|
|
process_data=process_data,
|
|
x_train=x_train,
|
|
x_test=x_test,
|
|
y_train=y_train,
|
|
y_test=y_test,
|
|
regr=regr,
|
|
)
|
|
|
|
assert result.params == sample_params
|
|
assert result.process_data == process_data
|
|
assert result.x_train.equals(x_train)
|
|
assert result.x_test.equals(x_test)
|
|
assert result.y_train.equals(y_train)
|
|
assert result.y_test.equals(y_test)
|
|
assert result.regr == regr
|
|
assert result.scaler_dict == scaler_dict
|
|
|
|
|
|
def test_train_model_result_optional_fields_default_none(sample_params, sample_dataframes):
|
|
"""Test that optional fields default to None."""
|
|
x_train, x_test, y_train, y_test = sample_dataframes
|
|
|
|
result = TrainModelResult(
|
|
params=sample_params,
|
|
process_data=MagicMock(),
|
|
x_train=x_train,
|
|
x_test=x_test,
|
|
y_train=y_train,
|
|
y_test=y_test,
|
|
regr=MagicMock(),
|
|
)
|
|
|
|
assert result.y_pred is None
|
|
assert result.mse_val is None
|
|
assert result.mae_val is None
|
|
assert result.r2_val is None
|
|
assert result.run_name is None
|
|
assert result.report_path is None
|
|
assert result.train_data_path is None
|
|
assert result.test_data_path is None
|
|
assert result.run_dir is None
|
|
|
|
|
|
def test_train_model_result_with_metrics(sample_params, sample_dataframes):
|
|
"""Test TrainModelResult with metrics populated."""
|
|
x_train, x_test, y_train, y_test = sample_dataframes
|
|
y_pred = pd.Series([41, 49])
|
|
|
|
result = TrainModelResult(
|
|
params=sample_params,
|
|
process_data=MagicMock(),
|
|
x_train=x_train,
|
|
x_test=x_test,
|
|
y_train=y_train,
|
|
y_test=y_test,
|
|
regr=MagicMock(),
|
|
y_pred=y_pred,
|
|
mse_val=1.5,
|
|
mae_val=1.2,
|
|
r2_val=0.95,
|
|
)
|
|
|
|
assert result.y_pred.equals(y_pred)
|
|
assert result.mse_val == 1.5
|
|
assert result.mae_val == 1.2
|
|
assert result.r2_val == 0.95
|
|
|
|
|
|
def test_train_model_result_with_artifact_paths(sample_params, sample_dataframes):
|
|
"""Test TrainModelResult with artifact paths populated."""
|
|
x_train, x_test, y_train, y_test = sample_dataframes
|
|
|
|
result = TrainModelResult(
|
|
params=sample_params,
|
|
process_data=MagicMock(),
|
|
x_train=x_train,
|
|
x_test=x_test,
|
|
y_train=y_train,
|
|
y_test=y_test,
|
|
regr=MagicMock(),
|
|
run_name='test-experiment-1',
|
|
report_path='/path/to/report.html',
|
|
train_data_path='/path/to/train_data.csv',
|
|
test_data_path='/path/to/test_data.csv',
|
|
run_dir='/path/to/run_dir',
|
|
)
|
|
|
|
assert result.run_name == 'test-experiment-1'
|
|
assert result.report_path == '/path/to/report.html'
|
|
assert result.train_data_path == '/path/to/train_data.csv'
|
|
assert result.test_data_path == '/path/to/test_data.csv'
|
|
assert result.run_dir == '/path/to/run_dir'
|
|
|
|
|
|
def test_train_model_result_is_dataclass(sample_params, sample_dataframes):
|
|
"""Test that TrainModelResult is a dataclass."""
|
|
x_train, x_test, y_train, y_test = sample_dataframes
|
|
|
|
result = TrainModelResult(
|
|
params=sample_params,
|
|
process_data=MagicMock(),
|
|
x_train=x_train,
|
|
x_test=x_test,
|
|
y_train=y_train,
|
|
y_test=y_test,
|
|
regr=MagicMock(),
|
|
scaler_dict={},
|
|
)
|
|
|
|
# Dataclasses have __dataclass_fields__ attribute
|
|
assert hasattr(result, '__dataclass_fields__')
|
|
assert 'params' in result.__dataclass_fields__
|
|
assert 'process_data' in result.__dataclass_fields__
|
|
assert 'x_train' in result.__dataclass_fields__
|
|
|
|
|
|
def test_train_model_result_field_count():
|
|
"""Test that TrainModelResult has exactly 19 fields."""
|
|
from dataclasses import fields
|
|
|
|
result_fields = fields(TrainModelResult)
|
|
assert len(result_fields) == 19
|
|
|
|
field_names = {f.name for f in result_fields}
|
|
expected_fields = {
|
|
'params',
|
|
'process_data',
|
|
'x_train',
|
|
'x_test',
|
|
'y_train',
|
|
'y_test',
|
|
'regr',
|
|
'y_pred',
|
|
'y_train_pred',
|
|
'mse_val',
|
|
'mae_val',
|
|
'r2_val',
|
|
'equation',
|
|
'equation_path',
|
|
'run_name',
|
|
'report_path',
|
|
'train_data_path',
|
|
'test_data_path',
|
|
'run_dir',
|
|
}
|
|
assert field_names == expected_fields
|
|
|
|
|
|
def test_train_model_result_complete_workflow(sample_params, sample_dataframes):
|
|
"""Test TrainModelResult through a complete workflow simulation."""
|
|
x_train, x_test, y_train, y_test = sample_dataframes
|
|
|
|
# Step 1: Create result after training
|
|
result = TrainModelResult(
|
|
params=sample_params,
|
|
process_data=MagicMock(),
|
|
x_train=x_train,
|
|
x_test=x_test,
|
|
y_train=y_train,
|
|
y_test=y_test,
|
|
regr=MagicMock(),
|
|
)
|
|
|
|
# Step 2: Add predictions and metrics
|
|
result.y_pred = pd.Series([41, 49])
|
|
result.mse_val = 1.5
|
|
result.mae_val = 1.2
|
|
result.r2_val = 0.95
|
|
|
|
# Step 3: Add artifact paths
|
|
result.run_name = 'test-experiment-1'
|
|
result.report_path = '/path/to/report.html'
|
|
result.train_data_path = '/path/to/train_data.csv'
|
|
result.test_data_path = '/path/to/test_data.csv'
|
|
result.run_dir = '/path/to/run_dir'
|
|
|
|
# Verify all fields are populated
|
|
assert result.y_pred is not None
|
|
assert result.mse_val == 1.5
|
|
assert result.mae_val == 1.2
|
|
assert result.r2_val == 0.95
|
|
assert result.run_name == 'test-experiment-1'
|
|
assert result.report_path == '/path/to/report.html'
|
|
assert result.train_data_path == '/path/to/train_data.csv'
|
|
assert result.test_data_path == '/path/to/test_data.csv'
|
|
assert result.run_dir == '/path/to/run_dir'
|