vitor-aignosi
424be007ef
Update dependencies and refactor MLFlow activities - Replaced direct GitHub dependencies in `requirements.txt` with specific versioned packages for `sientia_do` and `sientia_model`. - Refactored imports in `activities.py` to streamline the code structure. - Enhanced the `MLFlow` class in `mlflow.py` by introducing a method to resolve model aliases, improving flexibility in model lookups. - Simplified shutdown logic in `worker.py` for better readability. - Added new tests for MLFlow activities and improved existing test coverage for data handling and model retraining processes.
303 lines
9.4 KiB
Python
303 lines
9.4 KiB
Python
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
|
|
|
from pytest import mark
|
|
|
|
from laborious.activities.activities import Activities
|
|
from laborious.activities.api import API
|
|
from laborious.activities.gates import Gates
|
|
from laborious.activities.mlflow import MLFlow
|
|
from laborious.activities.model_metrics import ModelMetrics
|
|
from laborious.activities.opc import OPC
|
|
from laborious.activities.storage import Storage
|
|
|
|
|
|
@patch('laborious.activities.activities.Storage.__init__')
|
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
|
@patch('laborious.activities.activities.OPC.__init__')
|
|
@patch('laborious.activities.activities.Gates.__init__')
|
|
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
|
@patch('laborious.activities.activities.API.__init__')
|
|
@patch('laborious.activities.activities.MinioRepository')
|
|
@patch('laborious.activities.activities.MetricsController')
|
|
def test___init__(
|
|
mock_metrics_controller,
|
|
mock_minio_repository,
|
|
mock_api_init,
|
|
mock_model_metrics_init,
|
|
mock_gates_init,
|
|
mock_opc_init,
|
|
mock_mlflow_init,
|
|
mock_storage_init,
|
|
):
|
|
postgres_config = {
|
|
'host': 'localhost',
|
|
'port': 5432,
|
|
'user': 'postgres',
|
|
'password': 'postgres',
|
|
'dbname': 'postgres',
|
|
'min_connections': 1,
|
|
'max_connections': 10,
|
|
}
|
|
|
|
minio_config = {
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'default_bucket': 'test',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
}
|
|
|
|
mlflow_repository = MagicMock()
|
|
plugin_store = MagicMock()
|
|
|
|
opc_config = {
|
|
'bootstrap_servers': 'localhost:9092',
|
|
'polling_time': 1000,
|
|
'group_id': 'test-group',
|
|
}
|
|
|
|
pi_web_api_config = {
|
|
'base_url': 'https://test-pi-server.com',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'test_token',
|
|
}
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
activities = Activities(
|
|
postgres_config=postgres_config,
|
|
plugin_store=plugin_store,
|
|
minio_config=minio_config,
|
|
opc_config=opc_config,
|
|
pi_web_api_config=pi_web_api_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
mlflow_repository=mlflow_repository,
|
|
)
|
|
|
|
assert isinstance(activities, Activities)
|
|
assert isinstance(activities, Storage)
|
|
assert isinstance(activities, MLFlow)
|
|
assert isinstance(activities, OPC)
|
|
assert isinstance(activities, Gates)
|
|
assert isinstance(activities, ModelMetrics)
|
|
assert isinstance(activities, API)
|
|
|
|
mock_storage_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'],
|
|
retention_hours=minio_config['retention_hours'],
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_mlflow_init.assert_called_once_with(
|
|
ANY,
|
|
mlflow_repository=mlflow_repository,
|
|
plugin_store=plugin_store,
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_opc_init.assert_called_once_with(
|
|
ANY,
|
|
opc_servers=opc_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_gates_init.assert_called_once_with(
|
|
ANY,
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_model_metrics_init.assert_called_once_with(
|
|
ANY,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_api_init.assert_called_once_with(
|
|
ANY,
|
|
base_url=pi_web_api_config['base_url'],
|
|
auth_type=pi_web_api_config['auth_type'],
|
|
auth_token=pi_web_api_config['auth_token'],
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
endpoint=minio_config['endpoint_url'],
|
|
access_key=minio_config['access_key'],
|
|
secret_key=minio_config['secret_key'],
|
|
bucket=minio_config['default_bucket'],
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
secure=minio_config['secure'],
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.activities.Storage')
|
|
@patch('laborious.activities.activities.MLFlow')
|
|
@patch('laborious.activities.activities.OPC')
|
|
@patch('laborious.activities.activities.Gates')
|
|
@patch('laborious.activities.activities.ModelMetrics')
|
|
@patch('laborious.activities.activities.API')
|
|
@patch('laborious.activities.activities.MinioRepository')
|
|
async def test_shutdown(
|
|
_mock_minio_repository,
|
|
mock_api_init,
|
|
mock_model_metrics_init,
|
|
mock_gates_init,
|
|
mock_opc_init,
|
|
mock_mlflow_init,
|
|
mock_storage_init,
|
|
):
|
|
mock_opc_init.close = AsyncMock()
|
|
postgres_config = {
|
|
'host': 'localhost',
|
|
'port': 5432,
|
|
'user': 'postgres',
|
|
'password': 'postgres',
|
|
'dbname': 'postgres',
|
|
'min_connections': 1,
|
|
'max_connections': 10,
|
|
}
|
|
|
|
minio_config = {
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'default_bucket': 'test',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
}
|
|
|
|
mlflow_repository = MagicMock()
|
|
plugin_store = MagicMock()
|
|
|
|
opc_config = {
|
|
'bootstrap_servers': 'localhost:9092',
|
|
'polling_time': 1000,
|
|
'group_id': 'test-group',
|
|
}
|
|
|
|
pi_web_api_config = {
|
|
'base_url': 'https://test-pi-server.com',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'test_token',
|
|
}
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
activities = Activities(
|
|
postgres_config=postgres_config,
|
|
plugin_store=plugin_store,
|
|
minio_config=minio_config,
|
|
opc_config=opc_config,
|
|
pi_web_api_config=pi_web_api_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
mlflow_repository=mlflow_repository,
|
|
)
|
|
|
|
await activities.shutdown()
|
|
mock_opc_init.close.assert_called_once()
|
|
mock_storage_init.close.assert_called_once()
|
|
mock_mlflow_init.close.assert_called_once()
|
|
mock_gates_init.close.assert_called_once()
|
|
mock_model_metrics_init.close.assert_called_once()
|
|
mock_api_init.close.assert_called_once()
|
|
|
|
|
|
@patch('laborious.activities.activities.SientiaMLflowRepository')
|
|
@patch('laborious.activities.activities.build_mlflow_config')
|
|
@patch('laborious.activities.activities.Storage.__init__')
|
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
|
@patch('laborious.activities.activities.OPC.__init__')
|
|
@patch('laborious.activities.activities.Gates.__init__')
|
|
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
|
@patch('laborious.activities.activities.API.__init__')
|
|
@patch('laborious.activities.activities.MinioRepository')
|
|
@patch('laborious.activities.activities.MetricsController')
|
|
def test___init___builds_mlflow_repository_when_not_provided(
|
|
mock_metrics_controller,
|
|
mock_minio_repository,
|
|
_mock_api_init,
|
|
_mock_model_metrics_init,
|
|
_mock_gates_init,
|
|
_mock_opc_init,
|
|
_mock_mlflow_init,
|
|
_mock_storage_init,
|
|
mock_build_mlflow_config,
|
|
mock_mlflow_repository_cls,
|
|
):
|
|
postgres_config = {
|
|
'host': 'localhost',
|
|
'port': 5432,
|
|
'user': 'postgres',
|
|
'password': 'postgres',
|
|
'dbname': 'postgres',
|
|
'min_connections': 1,
|
|
'max_connections': 10,
|
|
}
|
|
minio_config = {
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'default_bucket': 'test',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
}
|
|
opc_config = {'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, 'group_id': 'test'}
|
|
pi_web_api_config = {'base_url': 'https://pi', 'auth_type': 'bearer', 'auth_token': 'token'}
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
plugin_store = MagicMock()
|
|
mock_build_mlflow_config.return_value = {
|
|
'url': 'http://mlflow:80',
|
|
'username': 'u',
|
|
'password': 'p',
|
|
}
|
|
|
|
Activities(
|
|
postgres_config=postgres_config,
|
|
plugin_store=plugin_store,
|
|
minio_config=minio_config,
|
|
opc_config=opc_config,
|
|
pi_web_api_config=pi_web_api_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
mock_build_mlflow_config.assert_called_once()
|
|
mock_mlflow_repository_cls.assert_called_once_with(
|
|
host='http://mlflow:80',
|
|
username='u',
|
|
password='p',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|