SIENTIAPDE-1241: refactor train_model workflow due to I/O errors.
This commit is contained in:
Bruno Domingues
@@ -1,535 +0,0 @@
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
from pytest import mark
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.activities.experiment_tracking import ExperimentTracking
|
||||
from model_manager.activities.mlflow import MLFlow
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
|
||||
@patch('model_manager.activities.activities.MLFlow.__init__')
|
||||
@patch('model_manager.activities.activities.MinIO.__init__')
|
||||
@patch('model_manager.activities.activities.Training.__init__')
|
||||
def test___init__(
|
||||
mock_training_init,
|
||||
mock_minio_init,
|
||||
mock_mlflow_init,
|
||||
mock_experiment_tracking_init,
|
||||
):
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||
|
||||
minio_config = {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'adaptive',
|
||||
'connect_timeout': 10,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
assert isinstance(activities, Activities)
|
||||
assert isinstance(activities, ExperimentTracking)
|
||||
assert isinstance(activities, MLFlow)
|
||||
assert isinstance(activities, Training)
|
||||
|
||||
mock_experiment_tracking_init.assert_called_once_with(
|
||||
ANY,
|
||||
host=postgres_config['host'],
|
||||
port=postgres_config['port'],
|
||||
user=postgres_config['user'],
|
||||
password=postgres_config['password'],
|
||||
dbname=postgres_config['dbname'],
|
||||
min_connections=postgres_config['min_connections'],
|
||||
max_connections=postgres_config['max_connections'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
mock_mlflow_init.assert_called_once_with(
|
||||
ANY,
|
||||
mlflow_host=mlflow_config['host'],
|
||||
mlflow_port=mlflow_config['port'],
|
||||
mlflow_username=mlflow_config['username'],
|
||||
mlflow_password=mlflow_config['password'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
mock_minio_init.assert_called_once_with(
|
||||
ANY,
|
||||
endpoint_url=minio_config['endpoint_url'],
|
||||
access_key=minio_config['access_key'],
|
||||
secret_key=minio_config['secret_key'],
|
||||
region=minio_config['region'],
|
||||
use_ssl=minio_config['use_ssl'],
|
||||
max_retry_attempts=minio_config['max_retry_attempts'],
|
||||
retry_mode=minio_config['retry_mode'],
|
||||
connect_timeout=minio_config['connect_timeout'],
|
||||
read_timeout=minio_config['read_timeout'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
mock_training_init.assert_called_once_with(
|
||||
ANY, logger=logger, notification_handler=notification_handler
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.activities.activities.ExperimentTracking', return_value=MagicMock())
|
||||
@patch('model_manager.activities.activities.MLFlow', return_value=MagicMock())
|
||||
async def test_shutdown(_mock_mlflow_init, mock_experiment_tracking_init):
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||
|
||||
minio_config = {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'adaptive',
|
||||
'connect_timeout': 10,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
await activities.shutdown()
|
||||
mock_experiment_tracking_init.close.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
|
||||
@patch('model_manager.activities.activities.MLFlow.__init__')
|
||||
@patch('model_manager.activities.activities.MinIO.__init__')
|
||||
@patch('model_manager.activities.activities.Training.__init__')
|
||||
def test___del___with_engine(
|
||||
mock_training_init,
|
||||
mock_minio_init,
|
||||
mock_mlflow_init,
|
||||
mock_experiment_tracking_init,
|
||||
):
|
||||
"""Test __del__ calls parent destructor when engine attribute exists."""
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||
|
||||
minio_config = {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'adaptive',
|
||||
'connect_timeout': 10,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
# Add engine attribute to simulate Postgres initialization
|
||||
activities.engine = MagicMock()
|
||||
|
||||
# Create a mock __del__ that will be detected by hasattr
|
||||
mock_parent_del = MagicMock()
|
||||
|
||||
# Patch both the class and the instance to ensure super().__del__ exists and is callable
|
||||
with patch.object(ExperimentTracking, '__del__', mock_parent_del, create=True):
|
||||
# Trigger __del__
|
||||
activities.__del__()
|
||||
|
||||
# Verify parent __del__ was called
|
||||
mock_parent_del.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
|
||||
@patch('model_manager.activities.activities.MLFlow.__init__')
|
||||
@patch('model_manager.activities.activities.MinIO.__init__')
|
||||
@patch('model_manager.activities.activities.Training.__init__')
|
||||
def test___del___without_engine(
|
||||
mock_training_init,
|
||||
mock_minio_init,
|
||||
mock_mlflow_init,
|
||||
mock_experiment_tracking_init,
|
||||
):
|
||||
"""Test __del__ does not call parent destructor when engine attribute is missing."""
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||
|
||||
minio_config = {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'adaptive',
|
||||
'connect_timeout': 10,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
# Ensure engine attribute does NOT exist
|
||||
if hasattr(activities, 'engine'):
|
||||
delattr(activities, 'engine')
|
||||
|
||||
# Mock super().__del__ to track if it's called
|
||||
with patch.object(ExperimentTracking, '__del__', MagicMock()) as mock_parent_del:
|
||||
# Trigger __del__
|
||||
activities.__del__()
|
||||
|
||||
# Verify parent __del__ was NOT called
|
||||
mock_parent_del.assert_not_called()
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
|
||||
@patch('model_manager.activities.activities.MLFlow.__init__')
|
||||
@patch('model_manager.activities.activities.MinIO.__init__')
|
||||
@patch('model_manager.activities.activities.Training.__init__')
|
||||
def test___del___handles_exception_gracefully(
|
||||
mock_training_init,
|
||||
mock_minio_init,
|
||||
mock_mlflow_init,
|
||||
mock_experiment_tracking_init,
|
||||
):
|
||||
"""Test __del__ handles exceptions from parent destructor gracefully."""
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||
|
||||
minio_config = {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'adaptive',
|
||||
'connect_timeout': 10,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
# Add engine attribute
|
||||
activities.engine = MagicMock()
|
||||
|
||||
# Mock super().__del__ to raise an exception
|
||||
mock_parent_del = MagicMock(side_effect=RuntimeError('Cleanup failed'))
|
||||
|
||||
with patch.object(ExperimentTracking, '__del__', mock_parent_del):
|
||||
# Trigger __del__ - should not raise exception
|
||||
try:
|
||||
activities.__del__()
|
||||
# Test passes if no exception is raised
|
||||
except Exception as e:
|
||||
# Test fails if exception propagates
|
||||
raise AssertionError(f'__del__ should not raise exception, but raised: {e}') from e
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
|
||||
@patch('model_manager.activities.activities.MLFlow.__init__')
|
||||
@patch('model_manager.activities.activities.MinIO.__init__')
|
||||
@patch('model_manager.activities.activities.Training.__init__')
|
||||
def test___del___when_parent_has_no_del(
|
||||
mock_training_init,
|
||||
mock_minio_init,
|
||||
mock_mlflow_init,
|
||||
mock_experiment_tracking_init,
|
||||
):
|
||||
"""Test __del__ handles case when parent class has no __del__ method."""
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||
|
||||
minio_config = {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'adaptive',
|
||||
'connect_timeout': 10,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
# Add engine attribute
|
||||
activities.engine = MagicMock()
|
||||
|
||||
# Remove __del__ from parent to simulate it not existing
|
||||
with patch.object(ExperimentTracking, '__del__', create=False):
|
||||
# Trigger __del__ - should not raise exception
|
||||
try:
|
||||
activities.__del__()
|
||||
# Test passes if no exception is raised
|
||||
except Exception as e:
|
||||
# Test fails if exception propagates
|
||||
raise AssertionError(
|
||||
f'__del__ should handle missing parent __del__, but raised: {e}'
|
||||
) from e
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
|
||||
@patch('model_manager.activities.activities.MLFlow.__init__')
|
||||
@patch('model_manager.activities.activities.MinIO.__init__')
|
||||
@patch('model_manager.activities.activities.Training.__init__')
|
||||
def test___del___calls_super_successfully(
|
||||
mock_training_init,
|
||||
mock_minio_init,
|
||||
mock_mlflow_init,
|
||||
mock_experiment_tracking_init,
|
||||
):
|
||||
"""Test __del__ successfully calls super().__del__() when it exists - covers line 118."""
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||
|
||||
minio_config = {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'adaptive',
|
||||
'connect_timeout': 10,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
# Mock all parent __init__ methods to return None
|
||||
mock_experiment_tracking_init.return_value = None
|
||||
mock_mlflow_init.return_value = None
|
||||
mock_minio_init.return_value = None
|
||||
mock_training_init.return_value = None
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
# Add engine attribute to simulate Postgres initialization
|
||||
activities.engine = MagicMock()
|
||||
|
||||
# Track if super().__del__() was actually called
|
||||
super_del_called = []
|
||||
|
||||
def mock_super_del(self):
|
||||
"""Mock parent __del__ that tracks when it's called."""
|
||||
super_del_called.append(True)
|
||||
|
||||
# Patch ExperimentTracking.__del__ to exist and be callable
|
||||
with patch.object(ExperimentTracking, '__del__', mock_super_del, create=True):
|
||||
# Trigger __del__ - this should execute line 118: super().__del__()
|
||||
activities.__del__()
|
||||
|
||||
# Verify that super().__del__() was actually called (line 118 executed)
|
||||
assert len(super_del_called) == 1, 'super().__del__() should have been called once'
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
|
||||
@patch('model_manager.activities.activities.MLFlow.__init__')
|
||||
@patch('model_manager.activities.activities.MinIO.__init__')
|
||||
@patch('model_manager.activities.activities.Training.__init__')
|
||||
def test___del___when_super_has_no_del_method(
|
||||
mock_training_init,
|
||||
mock_minio_init,
|
||||
mock_mlflow_init,
|
||||
mock_experiment_tracking_init,
|
||||
):
|
||||
"""Test __del__ handles case when hasattr(super(), '__del__') returns False - covers line 118 false branch."""
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||
|
||||
minio_config = {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'adaptive',
|
||||
'connect_timeout': 10,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
# Mock all __init__ methods to return None
|
||||
mock_experiment_tracking_init.return_value = None
|
||||
mock_mlflow_init.return_value = None
|
||||
mock_minio_init.return_value = None
|
||||
mock_training_init.return_value = None
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
# Add engine attribute to pass the first hasattr check (line 115)
|
||||
activities.engine = MagicMock()
|
||||
|
||||
# Create a mock class without __del__ method to simulate super() not having __del__
|
||||
class MockSuperWithoutDel:
|
||||
"""Mock class that explicitly does not have __del__ method."""
|
||||
|
||||
pass
|
||||
|
||||
# Patch super() to return an instance that doesn't have __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
|
||||
try:
|
||||
activities.__del__()
|
||||
# Test passes - the false branch of line 118 was executed without error
|
||||
except Exception as e:
|
||||
# Test fails if exception propagates
|
||||
raise AssertionError(
|
||||
f'__del__ should handle super() without __del__ method, but raised: {e}'
|
||||
) from e
|
||||
@@ -1,564 +0,0 @@
|
||||
"""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
|
||||
@@ -1,328 +0,0 @@
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from pytest import fixture, mark, raises
|
||||
|
||||
from model_manager.activities.minio import MinIO
|
||||
|
||||
|
||||
@patch('model_manager.activities.minio.boto3.client')
|
||||
def test___init__(mock_boto3_client):
|
||||
"""Test MinIO initialization with correct configuration."""
|
||||
mock_client = MagicMock()
|
||||
mock_boto3_client.return_value = mock_client
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
minio = MinIO(
|
||||
endpoint_url='http://localhost:9000',
|
||||
access_key='minioadmin',
|
||||
secret_key='minioadmin',
|
||||
region='us-east-1',
|
||||
use_ssl=False,
|
||||
max_retry_attempts=3,
|
||||
retry_mode='adaptive',
|
||||
connect_timeout=10,
|
||||
read_timeout=60,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
assert minio.endpoint_url == 'http://localhost:9000'
|
||||
assert minio.access_key == 'minioadmin'
|
||||
assert minio.secret_key == 'minioadmin'
|
||||
assert minio.region == 'us-east-1'
|
||||
assert minio.use_ssl is False
|
||||
assert minio.max_retry_attempts == 3
|
||||
assert minio.retry_mode == 'adaptive'
|
||||
assert minio.connect_timeout == 10
|
||||
assert minio.read_timeout == 60
|
||||
|
||||
# Verify boto3 client was created with correct parameters
|
||||
mock_boto3_client.assert_called_once()
|
||||
call_kwargs = mock_boto3_client.call_args[1]
|
||||
assert call_kwargs['endpoint_url'] == 'http://localhost:9000'
|
||||
assert call_kwargs['aws_access_key_id'] == 'minioadmin'
|
||||
assert call_kwargs['aws_secret_access_key'] == 'minioadmin'
|
||||
assert call_kwargs['use_ssl'] is False
|
||||
|
||||
|
||||
@patch('model_manager.activities.minio.boto3.client')
|
||||
def test___init___failure(mock_boto3_client):
|
||||
"""Test MinIO initialization failure handling."""
|
||||
mock_boto3_client.side_effect = Exception('Connection failed')
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
with raises(ConnectionError, match='Failed to initialize MinIO client'):
|
||||
MinIO(
|
||||
endpoint_url='http://localhost:9000',
|
||||
access_key='minioadmin',
|
||||
secret_key='minioadmin',
|
||||
region='us-east-1',
|
||||
use_ssl=False,
|
||||
max_retry_attempts=3,
|
||||
retry_mode='adaptive',
|
||||
connect_timeout=10,
|
||||
read_timeout=60,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('model_manager.activities.minio.boto3.client')
|
||||
def minio(mock_boto3_client):
|
||||
"""Fixture to create a MinIO instance for testing."""
|
||||
mock_client = MagicMock()
|
||||
mock_boto3_client.return_value = mock_client
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
minio_instance = MinIO(
|
||||
endpoint_url='http://localhost:9000',
|
||||
access_key='minioadmin',
|
||||
secret_key='minioadmin',
|
||||
region='us-east-1',
|
||||
use_ssl=False,
|
||||
max_retry_attempts=3,
|
||||
retry_mode='adaptive',
|
||||
connect_timeout=10,
|
||||
read_timeout=60,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
minio_instance.send_notification = MagicMock()
|
||||
minio_instance.minio_client = mock_client
|
||||
|
||||
return minio_instance
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'workflow_name': 'test_workflow',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_fetch_file_from_minio_success(minio):
|
||||
"""Test successful file fetch from MinIO."""
|
||||
# Arrange
|
||||
test_content = b'test file content'
|
||||
mock_response = {'Body': MagicMock()}
|
||||
mock_response['Body'].__enter__ = MagicMock(
|
||||
return_value=MagicMock(read=MagicMock(return_value=test_content))
|
||||
)
|
||||
mock_response['Body'].__exit__ = MagicMock(return_value=None)
|
||||
|
||||
minio.minio_client.get_object.return_value = mock_response
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'test-file.txt',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await minio.fetch_file_from_minio(input_data)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, BytesIO)
|
||||
result.seek(0)
|
||||
assert result.read() == test_content
|
||||
|
||||
minio.minio_client.get_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.txt')
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_fetch_file_from_minio_file_not_found(minio):
|
||||
"""Test file fetch when file doesn't exist."""
|
||||
# Arrange
|
||||
minio.minio_client.get_object.side_effect = Exception(
|
||||
'NoSuchKey: The specified key does not exist'
|
||||
)
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'nonexistent.txt',
|
||||
}
|
||||
|
||||
# Act & Assert
|
||||
with raises(OSError, match='Error fetching file from MinIO'):
|
||||
await minio.fetch_file_from_minio(input_data)
|
||||
|
||||
# Verify notification was sent
|
||||
minio.send_notification.assert_called_once()
|
||||
call_kwargs = minio.send_notification.call_args[1]
|
||||
assert call_kwargs['notification_id'] == 'FETCH_FILE_FROM_MINIO_ERROR'
|
||||
assert call_kwargs['block'] == 'fetch_file_from_minio'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_fetch_file_from_minio_network_error(minio):
|
||||
"""Test file fetch with network error."""
|
||||
# Arrange
|
||||
minio.minio_client.get_object.side_effect = Exception('Network timeout')
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'test-file.txt',
|
||||
}
|
||||
|
||||
# Act & Assert
|
||||
with raises(OSError, match='Error fetching file from MinIO'):
|
||||
await minio.fetch_file_from_minio(input_data)
|
||||
|
||||
minio.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_delete_file_from_minio_success(minio):
|
||||
"""Test successful file deletion from MinIO."""
|
||||
# Arrange
|
||||
minio.minio_client.delete_object.return_value = None
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'test-file.txt',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await minio.delete_file_from_minio(input_data)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
|
||||
minio.minio_client.delete_object.assert_called_once_with(
|
||||
Bucket='test-bucket', Key='test-file.txt'
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_delete_file_from_minio_idempotent(minio):
|
||||
"""Test that delete is idempotent (no error if file doesn't exist)."""
|
||||
# Arrange
|
||||
# MinIO delete_object is idempotent - no error if file doesn't exist
|
||||
minio.minio_client.delete_object.return_value = None
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'nonexistent.txt',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await minio.delete_file_from_minio(input_data)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
minio.minio_client.delete_object.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_delete_file_from_minio_access_denied(minio):
|
||||
"""Test file deletion with access denied error."""
|
||||
# Arrange
|
||||
minio.minio_client.delete_object.side_effect = Exception('AccessDenied: Access Denied')
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'test-file.txt',
|
||||
}
|
||||
|
||||
# Act & Assert
|
||||
with raises(OSError, match='Error deleting file from MinIO'):
|
||||
await minio.delete_file_from_minio(input_data)
|
||||
|
||||
# Verify notification was sent
|
||||
minio.send_notification.assert_called_once()
|
||||
call_kwargs = minio.send_notification.call_args[1]
|
||||
assert call_kwargs['notification_id'] == 'DELETE_FILE_FROM_MINIO_ERROR'
|
||||
assert call_kwargs['block'] == 'delete_file_from_minio'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_delete_file_from_minio_network_error(minio):
|
||||
"""Test file deletion with network error."""
|
||||
# Arrange
|
||||
minio.minio_client.delete_object.side_effect = Exception('Connection timeout')
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'test-file.txt',
|
||||
}
|
||||
|
||||
# Act & Assert
|
||||
with raises(OSError, match='Error deleting file from MinIO'):
|
||||
await minio.delete_file_from_minio(input_data)
|
||||
|
||||
minio.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_fetch_file_from_minio_large_file(minio):
|
||||
"""Test fetching a large file from MinIO."""
|
||||
# Arrange
|
||||
# Simulate a 10MB file
|
||||
large_content = b'x' * (10 * 1024 * 1024)
|
||||
mock_response = {'Body': MagicMock()}
|
||||
mock_response['Body'].__enter__ = MagicMock(
|
||||
return_value=MagicMock(read=MagicMock(return_value=large_content))
|
||||
)
|
||||
mock_response['Body'].__exit__ = MagicMock(return_value=None)
|
||||
|
||||
minio.minio_client.get_object.return_value = mock_response
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'large-file.bin',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await minio.fetch_file_from_minio(input_data)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, BytesIO)
|
||||
result.seek(0)
|
||||
assert len(result.read()) == 10 * 1024 * 1024
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_fetch_file_from_minio_empty_file(minio):
|
||||
"""Test fetching an empty file from MinIO."""
|
||||
# Arrange
|
||||
empty_content = b''
|
||||
mock_response = {'Body': MagicMock()}
|
||||
mock_response['Body'].__enter__ = MagicMock(
|
||||
return_value=MagicMock(read=MagicMock(return_value=empty_content))
|
||||
)
|
||||
mock_response['Body'].__exit__ = MagicMock(return_value=None)
|
||||
|
||||
minio.minio_client.get_object.return_value = mock_response
|
||||
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'empty-file.txt',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await minio.fetch_file_from_minio(input_data)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, BytesIO)
|
||||
result.seek(0)
|
||||
assert result.read() == b''
|
||||
@@ -1,447 +0,0 @@
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pytest import fixture, mark
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
from model_manager.activities.mlflow import MLFlow
|
||||
|
||||
|
||||
@patch('model_manager.activities.mlflow.MLFlowRepository')
|
||||
def test___init__(mock_mlflow_repository):
|
||||
mlflow = MLFlow(
|
||||
mlflow_host='http://localhost',
|
||||
mlflow_port=5000,
|
||||
mlflow_username='admin',
|
||||
mlflow_password='admin',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
)
|
||||
|
||||
assert mlflow.mlflow_host == 'http://localhost'
|
||||
assert mlflow.mlflow_port == 5000
|
||||
assert mlflow.mlflow_username == 'admin'
|
||||
assert mlflow.mlflow_password == 'admin'
|
||||
|
||||
mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('model_manager.activities.mlflow.MLFlowRepository')
|
||||
def mlflow(mock_mlflow_repository):
|
||||
mlflow = MLFlow(
|
||||
mlflow_host='http://localhost:5000',
|
||||
mlflow_port=5000,
|
||||
mlflow_username='admin',
|
||||
mlflow_password='admin',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
)
|
||||
|
||||
mlflow.send_notification = MagicMock()
|
||||
|
||||
return mlflow
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_save_model_success(mlflow):
|
||||
"""Test save_model successfully saves model and artifacts to MLflow."""
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
# Mock train result
|
||||
params = MagicMock(spec=TrainModelParams)
|
||||
params.experiment_name = 'test_experiment'
|
||||
|
||||
train_result = MagicMock(spec=TrainModelResult)
|
||||
train_result.params = params
|
||||
train_result.run_name = None # Will be set by get_next_run_name
|
||||
|
||||
# Mock repository methods
|
||||
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
||||
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
||||
mlflow.model_monitoring_repository.save_run.return_value = None
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'train_result': train_result,
|
||||
}
|
||||
|
||||
# Call the method
|
||||
response = await mlflow.save_model(input_data)
|
||||
|
||||
# Verify repository methods were called
|
||||
mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('test_experiment')
|
||||
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
|
||||
mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)
|
||||
|
||||
# Verify response - now returns TrainModelResult directly
|
||||
assert response == train_result
|
||||
assert response.run_name == 'test_experiment-1'
|
||||
assert train_result.run_name == 'test_experiment-1'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_save_model_get_next_run_name_error(mlflow):
|
||||
"""Test save_model handles error during get_next_run_name."""
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
params = MagicMock(spec=TrainModelParams)
|
||||
params.experiment_name = 'test_experiment'
|
||||
|
||||
train_result = MagicMock(spec=TrainModelResult)
|
||||
train_result.params = params
|
||||
|
||||
# Mock error in get_next_run_name
|
||||
mlflow.model_monitoring_repository.get_next_run_name.side_effect = Exception(
|
||||
'MLflow connection error'
|
||||
)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'train_result': train_result,
|
||||
}
|
||||
|
||||
# Call the method - should raise exception
|
||||
with pytest.raises(Exception, match='MLflow connection error'):
|
||||
await mlflow.save_model(input_data)
|
||||
|
||||
# Verify notification was sent
|
||||
mlflow.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='SAVE_MODEL_ERROR',
|
||||
message=ANY,
|
||||
block='save_model',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_save_model_generate_artifacts_error(mlflow):
|
||||
"""Test save_model handles error during generate_artifacts."""
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
params = MagicMock(spec=TrainModelParams)
|
||||
params.experiment_name = 'test_experiment'
|
||||
|
||||
train_result = MagicMock(spec=TrainModelResult)
|
||||
train_result.params = params
|
||||
|
||||
# Mock successful get_next_run_name but error in generate_artifacts
|
||||
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
||||
mlflow.model_monitoring_repository.generate_artifacts.side_effect = FileNotFoundError(
|
||||
'Reports directory does not exist'
|
||||
)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'train_result': train_result,
|
||||
}
|
||||
|
||||
# Call the method - should raise exception
|
||||
with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
|
||||
await mlflow.save_model(input_data)
|
||||
|
||||
# Verify notification was sent
|
||||
mlflow.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='SAVE_MODEL_ERROR',
|
||||
message=ANY,
|
||||
block='save_model',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_save_model_save_run_error(mlflow):
|
||||
"""Test save_model handles error during save_run."""
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
params = MagicMock(spec=TrainModelParams)
|
||||
params.experiment_name = 'test_experiment'
|
||||
|
||||
train_result = MagicMock(spec=TrainModelResult)
|
||||
train_result.params = params
|
||||
|
||||
# Mock successful get_next_run_name and generate_artifacts but error in save_run
|
||||
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
||||
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
||||
mlflow.model_monitoring_repository.save_run.side_effect = ValueError(
|
||||
'One or more metrics (MSE, R2, MAE) are None'
|
||||
)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'train_result': train_result,
|
||||
}
|
||||
|
||||
# Call the method - should raise exception
|
||||
with pytest.raises(ValueError, match=r'One or more metrics \(MSE, R2, MAE\) are None'):
|
||||
await mlflow.save_model(input_data)
|
||||
|
||||
# Verify notification was sent
|
||||
mlflow.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='SAVE_MODEL_ERROR',
|
||||
message=ANY,
|
||||
block='save_model',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_save_model_missing_metadata(mlflow):
|
||||
"""Test save_model handles missing metadata gracefully."""
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
params = MagicMock(spec=TrainModelParams)
|
||||
params.experiment_name = 'test_experiment'
|
||||
|
||||
train_result = MagicMock(spec=TrainModelResult)
|
||||
train_result.params = params
|
||||
|
||||
# Mock repository methods
|
||||
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
||||
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
||||
mlflow.model_monitoring_repository.save_run.return_value = None
|
||||
|
||||
# Input data without metadata
|
||||
input_data = {
|
||||
'train_result': train_result,
|
||||
}
|
||||
|
||||
# Call the method
|
||||
response = await mlflow.save_model(input_data)
|
||||
|
||||
# Verify it still works (metadata defaults to {}) - returns TrainModelResult directly
|
||||
assert response == train_result
|
||||
assert response.run_name == 'test_experiment-1'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_save_model_complete_flow(mlflow):
|
||||
"""Test save_model complete flow with all steps."""
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
params = MagicMock(spec=TrainModelParams)
|
||||
params.experiment_name = 'production_model'
|
||||
|
||||
train_result = MagicMock(spec=TrainModelResult)
|
||||
train_result.params = params
|
||||
train_result.run_name = None
|
||||
train_result.run_dir = None
|
||||
train_result.report_path = None
|
||||
|
||||
# Mock complete flow
|
||||
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'production_model-5'
|
||||
|
||||
# After generate_artifacts, paths should be set
|
||||
updated_result = MagicMock(spec=TrainModelResult)
|
||||
updated_result.params = params
|
||||
updated_result.run_name = 'production_model-5'
|
||||
updated_result.run_dir = '/reports/production_model-5_20231010'
|
||||
updated_result.report_path = '/reports/production_model-5_20231010/report.html'
|
||||
updated_result.train_data_path = '/reports/production_model-5_20231010/train_data.csv'
|
||||
updated_result.test_data_path = '/reports/production_model-5_20231010/test_data.csv'
|
||||
|
||||
mlflow.model_monitoring_repository.generate_artifacts.return_value = updated_result
|
||||
mlflow.model_monitoring_repository.save_run.return_value = None
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'train_result': train_result,
|
||||
}
|
||||
|
||||
# Call the method
|
||||
response = await mlflow.save_model(input_data)
|
||||
|
||||
# Verify complete flow
|
||||
mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('production_model')
|
||||
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
|
||||
mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)
|
||||
|
||||
# Verify response - returns TrainModelResult directly
|
||||
assert response == updated_result
|
||||
assert response.run_name == 'production_model-5'
|
||||
assert response.run_dir is not None
|
||||
assert response.report_path is not None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tests for cleanup_run_directory
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_run_directory_success(mlflow):
|
||||
"""Test successful cleanup of run directory."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'run_dir': 'test_run_dir',
|
||||
}
|
||||
|
||||
# Patch os and shutil inside the activity method
|
||||
with (
|
||||
patch('os.path.exists', return_value=True) as mock_exists,
|
||||
patch('shutil.rmtree') as mock_rmtree,
|
||||
):
|
||||
# Call the method
|
||||
await mlflow.cleanup_run_directory(input_data)
|
||||
|
||||
# Verify directory existence was checked
|
||||
mock_exists.assert_called_once_with('test_run_dir')
|
||||
|
||||
# Verify shutil.rmtree was called
|
||||
mock_rmtree.assert_called_once_with('test_run_dir')
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_run_directory_already_deleted(mlflow):
|
||||
"""Test cleanup when directory is already deleted (idempotent)."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'run_dir': 'already_deleted_dir',
|
||||
}
|
||||
|
||||
# Patch os and shutil inside the activity method
|
||||
with (
|
||||
patch('os.path.exists', return_value=False) as mock_exists,
|
||||
patch('shutil.rmtree') as mock_rmtree,
|
||||
):
|
||||
# Call the method - should not raise error
|
||||
await mlflow.cleanup_run_directory(input_data)
|
||||
|
||||
# Verify directory existence was checked
|
||||
mock_exists.assert_called_once_with('already_deleted_dir')
|
||||
|
||||
# Verify shutil.rmtree was NOT called
|
||||
mock_rmtree.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_run_directory_no_run_dir(mlflow):
|
||||
"""Test cleanup when no run_dir is provided."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
# No run_dir key
|
||||
}
|
||||
|
||||
# Call the method - should not raise error
|
||||
await mlflow.cleanup_run_directory(input_data)
|
||||
|
||||
# Should complete without errors
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_run_directory_empty_run_dir(mlflow):
|
||||
"""Test cleanup when run_dir is empty string."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'run_dir': '',
|
||||
}
|
||||
|
||||
# Call the method - should not raise error
|
||||
await mlflow.cleanup_run_directory(input_data)
|
||||
|
||||
# Should complete without errors
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_run_directory_none_run_dir(mlflow):
|
||||
"""Test cleanup when run_dir is None."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'run_dir': None,
|
||||
}
|
||||
|
||||
# Call the method - should not raise error
|
||||
await mlflow.cleanup_run_directory(input_data)
|
||||
|
||||
# Should complete without errors
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_run_directory_error(mlflow):
|
||||
"""Test cleanup handles errors correctly."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'run_dir': 'error_dir',
|
||||
}
|
||||
|
||||
# Patch os and shutil inside the activity method
|
||||
with (
|
||||
patch('os.path.exists', return_value=True),
|
||||
patch('shutil.rmtree', side_effect=PermissionError('Permission denied')),
|
||||
):
|
||||
# Call the method - should raise exception
|
||||
with pytest.raises(PermissionError, match='Permission denied'):
|
||||
await mlflow.cleanup_run_directory(input_data)
|
||||
|
||||
# Verify notification was sent
|
||||
mlflow.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
|
||||
message=ANY,
|
||||
block='cleanup_run_directory',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_run_directory_missing_metadata(mlflow):
|
||||
"""Test cleanup handles missing metadata gracefully."""
|
||||
input_data = {
|
||||
'run_dir': 'no_metadata_dir',
|
||||
}
|
||||
|
||||
# Patch os and shutil inside the activity method
|
||||
with patch('os.path.exists', return_value=True), patch('shutil.rmtree') as mock_rmtree:
|
||||
# Call the method
|
||||
await mlflow.cleanup_run_directory(input_data)
|
||||
|
||||
# Verify directory was deleted
|
||||
mock_rmtree.assert_called_once_with('no_metadata_dir')
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_run_directory_oserror(mlflow):
|
||||
"""Test cleanup handles OSError correctly."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'run_dir': 'os_error_dir',
|
||||
}
|
||||
|
||||
# Patch os and shutil inside the activity method
|
||||
with (
|
||||
patch('os.path.exists', return_value=True),
|
||||
patch('shutil.rmtree', side_effect=OSError('Directory not empty')),
|
||||
):
|
||||
# Call the method - should raise exception
|
||||
with pytest.raises(OSError, match='Directory not empty'):
|
||||
await mlflow.cleanup_run_directory(input_data)
|
||||
|
||||
# Verify notification was sent
|
||||
mlflow.send_notification.assert_called_once()
|
||||
call_args = mlflow.send_notification.call_args[1]
|
||||
assert call_args['notification_id'] == 'CLEANUP_RUN_DIRECTORY_ERROR'
|
||||
assert call_args['level'] == NotificationLevel.ERROR
|
||||
assert 'Directory not empty' in call_args['message']
|
||||
@@ -1,499 +0,0 @@
|
||||
"""Unit tests for Training activity."""
|
||||
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
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 - now returns TrainModelResult directly
|
||||
assert result == mock_final_result
|
||||
assert result.mse_val == 0.5
|
||||
assert result.mae_val == 0.3
|
||||
assert result.r2_val == 0.95
|
||||
|
||||
# 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,
|
||||
}
|
||||
|
||||
# Should raise ValueError
|
||||
with pytest.raises(ValueError, match='uploaded_file must be BytesIO'):
|
||||
await training.train_model(input_data)
|
||||
|
||||
# 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,
|
||||
}
|
||||
|
||||
# Should raise ValueError
|
||||
with pytest.raises(ValueError, match='Training data is empty'):
|
||||
await training.train_model(input_data)
|
||||
|
||||
# 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,
|
||||
}
|
||||
|
||||
# Should raise Exception
|
||||
with pytest.raises(Exception, match='Database connection failed'):
|
||||
await training.train_model(input_data)
|
||||
|
||||
# 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_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,
|
||||
}
|
||||
|
||||
# Should raise Exception
|
||||
with pytest.raises(Exception, match='Metric calculation failed'):
|
||||
await training.train_model(input_data)
|
||||
|
||||
# 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
|
||||
}
|
||||
|
||||
# Should raise ValueError
|
||||
with pytest.raises(ValueError, match='train_params must be TrainModelParams.*dict'):
|
||||
await training.train_model(input_data)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tests for validate_train_params
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_validate_train_params_success():
|
||||
"""Test successful validation of training parameters."""
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
training.info = MagicMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'experiment_run_id': 456,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1', 'feature2', 'price'], # target_variable must be in list
|
||||
'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, 'price': 0.0},
|
||||
'upp_lim': {'feature1': 100.0, 'feature2': 100.0, 'price': 1000.0},
|
||||
'window': 10,
|
||||
'experiment_name': 'test_experiment',
|
||||
'removed_intervals': [],
|
||||
}
|
||||
|
||||
result = await training.validate_train_params(input_data)
|
||||
|
||||
assert isinstance(result, TrainModelParams)
|
||||
assert result.experiment_run_id == 456
|
||||
assert result.target_variable == 'price'
|
||||
assert result.variable_columns == ['feature1', 'feature2', 'price']
|
||||
assert result.train_size == 80
|
||||
assert result.experiment_name == 'test_experiment'
|
||||
assert training.info.call_count == 2
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_validate_train_params_missing_required_field():
|
||||
"""Test validation fails when required field is missing."""
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
training.info = MagicMock()
|
||||
training.error = MagicMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'experiment_run_id': 456,
|
||||
'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',
|
||||
'removed_intervals': [],
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match='target_variable'):
|
||||
await training.validate_train_params(input_data)
|
||||
|
||||
notification_handler.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_validate_train_params_invalid_type():
|
||||
"""Test validation fails when field has invalid type."""
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
training.info = MagicMock()
|
||||
training.error = MagicMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-invalid'},
|
||||
'experiment_run_id': 456,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1'],
|
||||
'train_size': 'invalid',
|
||||
'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',
|
||||
'removed_intervals': [],
|
||||
}
|
||||
|
||||
with pytest.raises((ValueError, TypeError)):
|
||||
await training.validate_train_params(input_data)
|
||||
|
||||
notification_handler.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_validate_train_params_empty_input():
|
||||
"""Test validation fails with empty input."""
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
training.info = MagicMock()
|
||||
training.error = MagicMock()
|
||||
|
||||
input_data = {'metadata': {}}
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
await training.validate_train_params(input_data)
|
||||
|
||||
notification_handler.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_validate_train_params_without_metadata():
|
||||
"""Test validation works even without metadata key."""
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
training.info = MagicMock()
|
||||
|
||||
input_data = {
|
||||
'experiment_run_id': 789,
|
||||
'target_variable': 'temperature',
|
||||
'variable_columns': ['sensor1', 'temperature'], # target_variable must be in list
|
||||
'train_size': 75,
|
||||
'shuffle': False,
|
||||
'use_scaler': True,
|
||||
'include_ar': True,
|
||||
'bucket_name': 'sensors',
|
||||
'file_name': 'data.csv',
|
||||
'line_separator': '\n',
|
||||
'decimal_separator': '.',
|
||||
'lag_train': 2,
|
||||
'lag_val': 2,
|
||||
'rem_static_win': True,
|
||||
'low_lim': {'sensor1': -50.0, 'temperature': -50.0},
|
||||
'upp_lim': {'sensor1': 150.0, 'temperature': 150.0},
|
||||
'window': 20,
|
||||
'experiment_name': 'sensor_experiment',
|
||||
'removed_intervals': [],
|
||||
}
|
||||
|
||||
result = await training.validate_train_params(input_data)
|
||||
|
||||
assert isinstance(result, TrainModelParams)
|
||||
assert result.experiment_run_id == 789
|
||||
assert result.target_variable == 'temperature'
|
||||
assert result.experiment_name == 'sensor_experiment'
|
||||
Reference in New Issue
Block a user