from unittest.mock import ANY, MagicMock, patch from pytest import mark from laborious.activities.activities import Activities from laborious.activities.gates import Gates from laborious.activities.mlflow import MLFlow 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__') def test___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', 'region_name': 'us-east-1', 'default_bucket': 'test', } 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, minio_config=minio_config, opc_config=opc_config, logger=logger, notification_handler=notification_handler, ) assert isinstance(activities, Activities) assert isinstance(activities, Storage) assert isinstance(activities, MLFlow) assert isinstance(activities, OPC) assert isinstance(activities, Gates) 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'], minio_config=minio_config, 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'], minio_config=minio_config, 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('laborious.activities.activities.Storage', return_value=MagicMock()) @patch('laborious.activities.activities.MLFlow', return_value=MagicMock()) @patch('laborious.activities.activities.OPC', return_value=MagicMock()) async def test_shutdown(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', 'region_name': 'us-east-1', 'default_bucket': 'test', } 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, minio_config=minio_config, opc_config=opc_config, logger=logger, notification_handler=notification_handler, ) await activities.shutdown() mock_opc_init.shutdown.assert_called_once() mock_storage_init.close.assert_called_once()