from unittest.mock import patch, MagicMock, ANY from scouter.activities.activities import Activities from scouter.activities.postgres import Postgres from scouter.activities.redis import Redis from scouter.activities.kafka import Kafka from scouter.activities.gates import Gates from pytest import mark @patch('scouter.activities.activities.Postgres.__init__') @patch('scouter.activities.activities.Redis.__init__') @patch('scouter.activities.activities.Kafka.__init__') @patch('scouter.activities.activities.Gates.__init__') def test___init__(mock_gates_init, mock_kafka_init, mock_redis_init, mock_postgres_init): postgres_config = { 'host': 'localhost', 'port': 5432, 'user': 'postgres', 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, 'max_connections': 10 } redis_config = { 'host': 'localhost', 'port': 6379, 'username': 'redis', 'password': 'redis' } kafka_config = { 'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, 'group_id': 'test-group' } logger = MagicMock() notification_handler = MagicMock() activities = Activities( postgres_config=postgres_config, redis_config=redis_config, kafka_config=kafka_config, logger=logger, notification_handler=notification_handler ) assert isinstance(activities, Activities) assert isinstance(activities, Postgres) assert isinstance(activities, Redis) assert isinstance(activities, Kafka) 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_redis_init.assert_called_once_with( ANY, host=redis_config['host'], port=redis_config['port'], username=redis_config['username'], password=redis_config['password'], logger=logger, notification_handler=notification_handler ) mock_kafka_init.assert_called_once_with( ANY, bootstrap_servers=kafka_config['bootstrap_servers'], polling_time=kafka_config['polling_time'], group_id=kafka_config['group_id'], logger=logger, notification_handler=notification_handler ) mock_gates_init.assert_called_once_with( ANY, logger=logger, notification_handler=notification_handler ) @mark.asyncio @patch('scouter.activities.activities.Postgres.__init__') @patch('scouter.activities.activities.Redis.__init__') @patch('scouter.activities.activities.Kafka.__init__') async def test_prepare_activity(_mock_kafka_init, _mock_redis_init, _mock_postgres_init): postgres_config = { 'host': 'localhost', 'port': 5432, 'user': 'postgres', 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, 'max_connections': 10 } redis_config = { 'host': 'localhost', 'port': 6379, 'username': 'redis', 'password': 'redis' } kafka_config = { 'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, 'group_id': 'test-group' } logger = MagicMock() notification_handler = MagicMock() activities = Activities( postgres_config=postgres_config, redis_config=redis_config, kafka_config=kafka_config, logger=logger, notification_handler=notification_handler ) input_data = { 'workflow_name': 'test-workflow-name', 'schedule_name': 'test-schedule-name', 'model_name': 'test-model-name', 'model_id': 'test-model-id' } await activities.prepare_activity(input_data) assert activities.notification_handler.base_notification.pipeline_name == input_data[ 'workflow_name'] assert activities.notification_handler.base_notification.schedule_name == input_data[ 'schedule_name'] assert activities.notification_handler.base_notification.model_name == input_data[ 'model_name'] assert activities.notification_handler.base_notification.model_id == input_data[ 'model_id']