aba4a3f1a5
This commit renames the 'laborious' package to 'model_manager' across the entire project. This includes renaming directories, modules, references in code, configuration files, and documentation to reflect the new package name. This change improves clarity and consistency within the project.
136 lines
4.0 KiB
Python
136 lines
4.0 KiB
Python
from pytest import mark
|
|
from unittest.mock import patch, MagicMock, ANY
|
|
from sientia_do.temporal.activities.postgres import Postgres
|
|
from model_manager.activities.activities import Activities
|
|
from model_manager.activities.mlflow import MLFlow
|
|
from model_manager.activities.gates import Gates
|
|
from model_manager.activities.opc import OPC
|
|
|
|
|
|
@patch('model_manager.activities.activities.Postgres.__init__')
|
|
@patch('model_manager.activities.activities.MLFlow.__init__')
|
|
@patch('model_manager.activities.activities.OPC.__init__')
|
|
@patch('model_manager.activities.activities.Gates.__init__')
|
|
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_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'
|
|
}
|
|
|
|
opc_config = {
|
|
'bootstrap_servers': 'localhost:9092',
|
|
'polling_time': 1000,
|
|
'group_id': 'test-group'
|
|
}
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
activities = Activities(
|
|
postgres_config=postgres_config,
|
|
mlflow_config=mlflow_config,
|
|
opc_config=opc_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
assert isinstance(activities, Activities)
|
|
assert isinstance(activities, Postgres)
|
|
assert isinstance(activities, MLFlow)
|
|
assert isinstance(activities, OPC)
|
|
assert isinstance(activities, Gates)
|
|
|
|
mock_postgres_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_opc_init.assert_called_once_with(
|
|
ANY,
|
|
opc_servers=opc_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
mock_gates_init.assert_called_once_with(
|
|
ANY,
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.activities.Postgres', return_value=MagicMock())
|
|
@patch('model_manager.activities.activities.MLFlow', return_value=MagicMock())
|
|
@patch('model_manager.activities.activities.OPC', return_value=MagicMock())
|
|
async def test_shutdown(mock_opc_init,
|
|
_mock_mlflow_init, mock_postgres_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'
|
|
}
|
|
|
|
opc_config = {
|
|
'bootstrap_servers': 'localhost:9092',
|
|
'polling_time': 1000,
|
|
'group_id': 'test-group'
|
|
}
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
activities = Activities(
|
|
postgres_config=postgres_config,
|
|
mlflow_config=mlflow_config,
|
|
opc_config=opc_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
await activities.shutdown()
|
|
mock_opc_init.shutdown.assert_called_once()
|
|
mock_postgres_init.close.assert_called_once()
|