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

"""Unit tests for ExperimentTracking activity."""

from unittest.mock import AsyncMock, MagicMock, patch

from pytest import mark, raises

from model_manager.activities.experiment_tracking import ExperimentTracking, UpdateType


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_status_success(mock_postgres_init):
    """Test successful status update."""
    mock_postgres_init.return_value = None

    # Create instance
    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    # Mock methods
    tracking.info = MagicMock()
    tracking.execute_query = AsyncMock(return_value={'rowcount': 1})

    # Test data
    input_data = {
        'metadata': {'workflow_id': 'test-123'},
        'experiment_run_id': 456,
        'update_type': UpdateType.STATUS,
        'status': 'TRAINING_SUCCESS',
    }

    # Execute
    await tracking.update_experiment_run(input_data)

    # Assertions
    tracking.info.assert_called()
    tracking.execute_query.assert_called_once()
    call_args = tracking.execute_query.call_args
    assert 'UPDATE experiment_run' in call_args[0][0]
    assert 'SET status = %s' in call_args[0][0]


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_status_with_error_success(mock_postgres_init):
    """Test successful status update with error message."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.info = MagicMock()
    tracking.execute_query = AsyncMock(return_value={'rowcount': 1})

    input_data = {
        'metadata': {'workflow_id': 'test-123'},
        'experiment_run_id': 456,
        'update_type': UpdateType.STATUS_WITH_ERROR,
        'status': 'TRAINING_ERROR',
        'error_message': 'Model training failed due to insufficient data',
    }

    await tracking.update_experiment_run(input_data)

    tracking.info.assert_called()
    tracking.execute_query.assert_called_once()
    call_args = tracking.execute_query.call_args
    assert 'UPDATE experiment_run' in call_args[0][0]
    assert 'SET status = %s, error_message = %s' in call_args[0][0]


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_error_message_truncation(mock_postgres_init):
    """Test that error messages longer than 1024 chars are truncated."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.info = MagicMock()
    tracking.execute_query = AsyncMock(return_value={'rowcount': 1})

    # Create error message longer than 1024 characters
    long_error = 'A' * 2000

    input_data = {
        'metadata': {},
        'experiment_run_id': 456,
        'update_type': UpdateType.STATUS_WITH_ERROR,
        'status': 'TRAINING_ERROR',
        'error_message': long_error,
    }

    await tracking.update_experiment_run(input_data)

    # Check that error message was truncated to 1024 chars
    call_args = tracking.execute_query.call_args
    query_params = call_args[0][1]
    assert len(query_params[1]) == 1024  # error_message is second param


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_model_saved_success(mock_postgres_init):
    """Test successful model saved update."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.info = MagicMock()
    tracking.execute_query = AsyncMock(return_value={'rowcount': 1})

    input_data = {
        'metadata': {},
        'experiment_run_id': 456,
        'update_type': UpdateType.MODEL_SAVED,
        'run_name': 'experiment-model-123',
        'status': 'MLFLOW_SENT',
    }

    await tracking.update_experiment_run(input_data)

    tracking.info.assert_called()
    tracking.execute_query.assert_called_once()
    call_args = tracking.execute_query.call_args
    assert 'UPDATE experiment_run' in call_args[0][0]
    assert 'SET run_name = %s, status = %s' in call_args[0][0]


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_missing_status_raises_error(mock_postgres_init):
    """Test that missing status parameter raises ValueError."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.send_notification = MagicMock()
    tracking.error = MagicMock()

    input_data = {
        'metadata': {},
        'experiment_run_id': 456,
        'update_type': UpdateType.STATUS,
        # Missing 'status' parameter
    }

    with raises(RuntimeError, match='Error updating experiment run'):
        await tracking.update_experiment_run(input_data)

    tracking.send_notification.assert_called_once()


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_missing_error_message_raises_error(mock_postgres_init):
    """Test that missing error_message parameter raises ValueError."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.send_notification = MagicMock()
    tracking.error = MagicMock()

    input_data = {
        'metadata': {},
        'experiment_run_id': 456,
        'update_type': UpdateType.STATUS_WITH_ERROR,
        'status': 'TRAINING_ERROR',
        # Missing 'error_message' parameter
    }

    with raises(RuntimeError, match='Error updating experiment run'):
        await tracking.update_experiment_run(input_data)

    tracking.send_notification.assert_called_once()


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_missing_run_name_raises_error(mock_postgres_init):
    """Test that missing run_name parameter raises ValueError."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.send_notification = MagicMock()
    tracking.error = MagicMock()

    input_data = {
        'metadata': {},
        'experiment_run_id': 456,
        'update_type': UpdateType.MODEL_SAVED,
        # Missing 'run_name' parameter
    }

    with raises(RuntimeError, match='Error updating experiment run'):
        await tracking.update_experiment_run(input_data)

    tracking.send_notification.assert_called_once()


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_invalid_update_type_raises_error(mock_postgres_init):
    """Test that invalid update_type raises ValueError."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.send_notification = MagicMock()
    tracking.error = MagicMock()

    input_data = {
        'metadata': {},
        'experiment_run_id': 456,
        'update_type': 'INVALID_TYPE',
        'status': 'TRAINING_SUCCESS',
    }

    with raises(RuntimeError, match='Error updating experiment run'):
        await tracking.update_experiment_run(input_data)

    tracking.send_notification.assert_called_once()


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_no_rows_updated_raises_error(mock_postgres_init):
    """Test that zero rows updated raises ValueError."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.info = MagicMock()
    tracking.execute_query = AsyncMock(return_value={'rowcount': 0})
    tracking.send_notification = MagicMock()
    tracking.error = MagicMock()

    input_data = {
        'metadata': {},
        'experiment_run_id': 999,  # Non-existent ID
        'update_type': UpdateType.STATUS,
        'status': 'TRAINING_SUCCESS',
    }

    with raises(RuntimeError, match='Error updating experiment run'):
        await tracking.update_experiment_run(input_data)

    tracking.send_notification.assert_called_once()


@mark.asyncio
@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
async def test_update_experiment_run_sends_notification_on_error(mock_postgres_init):
    """Test that notification is sent when update fails."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    tracking.info = MagicMock()
    tracking.execute_query = AsyncMock(side_effect=Exception('Database error'))
    tracking.send_notification = MagicMock()
    tracking.error = MagicMock()

    input_data = {
        'metadata': {'workflow_id': 'test-123'},
        'experiment_run_id': 456,
        'update_type': UpdateType.STATUS,
        'status': 'TRAINING_SUCCESS',
    }

    with raises(RuntimeError):
        await tracking.update_experiment_run(input_data)

    # Verify notification was sent
    tracking.send_notification.assert_called_once()
    call_args = tracking.send_notification.call_args
    assert call_args[1]['notification_id'] == 'UPDATE_EXPERIMENT_RUN_ERROR'
    assert call_args[1]['metadata'] == {'workflow_id': 'test-123'}


def test_update_type_enum_values():
    """Test UpdateType enum has correct values."""
    assert UpdateType.STATUS == 'status'
    assert UpdateType.STATUS_WITH_ERROR == 'status_with_error'
    assert UpdateType.MODEL_SAVED == 'model_saved'


# Tests for __del__ method - 100% coverage


@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
def test___del___with_engine_and_parent_del_exists(mock_postgres_init):
    """Test __del__ calls parent destructor when engine exists and parent has __del__."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    # Add engine attribute to simulate Postgres initialization
    tracking.engine = MagicMock()

    # Track if parent __del__ was called
    parent_del_called = []

    def mock_parent_del(self):
        """Mock parent __del__ that tracks when it's called."""
        parent_del_called.append(True)

    # Patch parent class to have __del__ method
    with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True):
        # Trigger __del__ - should call parent __del__ (line 103)
        tracking.__del__()

        # Verify parent __del__ was called (line 103 executed)
        assert len(parent_del_called) == 1, 'Parent __del__ should have been called once'


@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
def test___del___without_engine(mock_postgres_init):
    """Test __del__ does not call parent destructor when engine attribute is missing."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    # Ensure engine attribute does not exist
    if hasattr(tracking, 'engine'):
        delattr(tracking, 'engine')

    # Mock parent __del__ to track if it's called
    mock_parent_del = MagicMock()

    with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True):
        # Trigger __del__ - should NOT call parent __del__ (line 100 is False)
        tracking.__del__()

        # Verify parent __del__ was NOT called
        mock_parent_del.assert_not_called()


@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
def test___del___when_parent_has_no_del_method(mock_postgres_init):
    """Test __del__ handles case when parent class has no __del__ method - covers line 102 false branch."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    # Add engine attribute to pass the first hasattr check (line 100)
    tracking.engine = MagicMock()

    # Create a mock class without __del__ method
    class MockSuperWithoutDel:
        """Mock class that explicitly does not have __del__ method."""

        pass

    # Patch super() to return an instance without __del__
    mock_super_instance = MockSuperWithoutDel()

    with patch('builtins.super', return_value=mock_super_instance):
        # Trigger __del__ - should handle the case when hasattr(super(), '__del__') is False (line 102)
        try:
            tracking.__del__()
            # Test passes - the false branch of line 102 was executed without error
        except Exception as e:
            raise AssertionError(
                f'__del__ should handle super() without __del__ method, but raised: {e}'
            ) from e


@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
def test___del___handles_exception_from_parent_del(mock_postgres_init):
    """Test __del__ handles exceptions from parent destructor gracefully - covers line 104."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    # Add engine attribute
    tracking.engine = MagicMock()

    # Mock parent __del__ to raise an exception
    mock_parent_del = MagicMock(side_effect=RuntimeError('Cleanup failed'))

    with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True):
        # Trigger __del__ - should catch exception and not propagate it (line 104-106)
        try:
            tracking.__del__()
            # Test passes if no exception is raised
        except Exception as e:
            raise AssertionError(
                f'__del__ should handle exceptions gracefully, but raised: {e}'
            ) from e


@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
def test___del___handles_attribute_error_from_parent_del(mock_postgres_init):
    """Test __del__ handles AttributeError from parent destructor - covers line 104 exception handling."""
    mock_postgres_init.return_value = None

    tracking = ExperimentTracking(
        host='localhost',
        port=5432,
        user='test',
        password='test',
        dbname='test',
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    # Add engine attribute
    tracking.engine = MagicMock()

    # Mock parent __del__ to raise AttributeError
    mock_parent_del = MagicMock(side_effect=AttributeError('engine not found'))

    with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True):
        # Trigger __del__ - should catch AttributeError and not propagate it
        try:
            tracking.__del__()
            # Test passes if no exception is raised
        except Exception as e:
            raise AssertionError(
                f'__del__ should handle AttributeError gracefully, but raised: {e}'
            ) from e