sientia-dataops-model-manager/tests/activities/test_training.py

"""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_params import TrainModelParams
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 = 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,
    }

    # 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)
    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': train_params,
    }

    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')
    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': train_params,
    }

    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')
    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': train_params,
    }

    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')
    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': train_params,
    }

    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()


@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()