from pytest import mark from unittest.mock import patch, MagicMock, ANY from sientia_do.temporal.activities.postgres import Postgres from laborious.activities.activities import Activities from laborious.activities.mlflow import MLFlow from laborious.activities.gates import Gates from laborious.activities.opc import OPC @patch('laborious.activities.activities.Postgres.__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_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('laborious.activities.activities.Postgres', 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_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()