from unittest.mock import ANY, AsyncMock, MagicMock, call, patch from pytest import fixture, mark from ingestor.ingestor import Ingestor @patch('ingestor.ingestor.getenv') @patch('ingestor.ingestor.NotificationHandler') def test___init__(notification_handler, getenv): getenv.side_effect = [ 'localhost:9092,localhost:35', # KAFKA_SERVERS 'true', # EXPORT_TO_KAFKA 'localhost', # REDIS_HOST '63790', # REDIS_PORT 'user', # REDIS_USERNAME 'password', # REDIS_PASSWORD '100', # LEASE_TTL '200', # HEARTBEAT_TTL 'localhost', # HOSTNAME '50', # POLL_INTERVAL 'localhost:27017', # MONGODB_URL 'sientia', # MONGODB_USERNAME 'sientia', # MONGODB_PASSWORD 'sientia', # MONGODB_DATABASE ] ingestor = Ingestor() getenv.assert_any_call('KAFKA_SERVERS', 'localhost:9092') getenv.assert_any_call('REDIS_HOST', 'localhost') getenv.assert_any_call('REDIS_PORT', '6379') getenv.assert_any_call('REDIS_USERNAME', None) getenv.assert_any_call('REDIS_PASSWORD', None) getenv.assert_any_call('LEASE_TTL', '10') getenv.assert_any_call('HEARTBEAT_TTL', '20') getenv.assert_any_call('HOSTNAME', 'localhost') getenv.assert_any_call('POLL_INTERVAL', '5') assert ingestor.kafka_servers == ['localhost:9092', 'localhost:35'] assert ingestor.redis_host == 'localhost' assert ingestor.redis_port == 63790 assert ingestor.redis_username == 'user' assert ingestor.redis_password == 'password' assert ingestor.lease_ttl == 100 assert ingestor.heartbeat_ttl == 200 assert ingestor.pod_id == 'localhost' assert ingestor.poll_interval == 50 assert ingestor.metadata == { 'model_id': '-', 'model_name': '-', 'workflow_name': 'opc_ingestor', 'schema_name': 'opc_ingestor', 'pod_id': 'localhost', } notification_handler.assert_called_once_with( connection_string='mongodb://sientia:sientia@localhost:27017', database='sientia', logger=ingestor.logger, project_name='opc_ingestor', ) @fixture @patch('ingestor.ingestor.getenv') @patch('ingestor.ingestor.NotificationHandler') def ingestor(_notification_handler, _getenv): ing = Ingestor() ing.logger = MagicMock() return ing @fixture def ingestor_manager_started(ingestor): ingestor.ingestor_manager = AsyncMock( initialize_opc_from_config=AsyncMock(), shutdown=AsyncMock(), update_opc_servers=AsyncMock(), subscribe_to_tags=AsyncMock(), unsubscribe_slot=AsyncMock(), ) return ingestor @mark.asyncio async def test_shutdown(ingestor_manager_started): await ingestor_manager_started.shutdown() ingestor_manager_started.ingestor_manager.shutdown.assert_called_once() @mark.asyncio async def test_handle_acquired_tags_not_acquired(ingestor_manager_started): await ingestor_manager_started.handle_acquired_tags([]) ingestor_manager_started.logger.warning.assert_called_once_with('No slots available') ingestor_manager_started.ingestor_manager.update_opc_servers.assert_not_called() ingestor_manager_started.ingestor_manager.subscribe_to_tags.assert_not_called() @mark.asyncio async def test_handle_acquired_tags_success(ingestor_manager_started): await ingestor_manager_started.handle_acquired_tags(['tag1', 'tag2']) ingestor_manager_started.logger.warning.assert_not_called() ingestor_manager_started.ingestor_manager.update_opc_servers.assert_called_once() ingestor_manager_started.ingestor_manager.subscribe_to_tags.assert_called_once_with( ['tag1', 'tag2'] ) @patch('ingestor.ingestor.IngestorManager') @mark.asyncio async def test_prepare_ingestor(ingestor_manager_mock, ingestor): ingestor_manager = ingestor_manager_mock.return_value ingestor_manager.get_slot_leases.return_value = True ingestor_manager.declare_active = AsyncMock() ingestor_manager.get_slot_leases = AsyncMock() ingestor.handle_acquired_tags = AsyncMock() await ingestor.prepare_ingestor() ingestor_manager_mock.assert_called_once_with( kafka_servers=','.join(ingestor.kafka_servers), redis_data={ 'host': ingestor.redis_host, 'port': ingestor.redis_port, 'username': ingestor.redis_username, 'password': ingestor.redis_password, }, lease_ttl=ingestor.lease_ttl, heartbeat_ttl=ingestor.heartbeat_ttl, poll_interval=ingestor.poll_interval, mongo_connection_string=ingestor.mongo_connection_string, mongo_database=ingestor.mongo_database, metadata=ingestor.metadata, logger=ingestor.logger, notification_handler=ingestor.notification_handler, export_to_kafka=ingestor.export_to_kafka, metrics_controller=ingestor.metrics_controller, ) ingestor_manager.declare_active.assert_called_once() ingestor_manager.get_slot_leases.assert_called_once() ingestor.handle_acquired_tags.assert_called_once_with( ingestor_manager.get_slot_leases.return_value ) def test_manage_slots_has_slots(ingestor_manager_started): ingestor_manager_started.handle_acquired_tags = MagicMock() ingestor_manager_started.ingestor_manager.managed_tags = True ingestor_manager_started.manage_no_slots(5) ingestor_manager_started.ingestor_manager.get_slot_leases.assert_not_called() ingestor_manager_started.ingestor_manager.handle_acquired_tags.assert_not_called() def test_manage_no_slots_has_slots_none_available(ingestor_manager_started): ingestor_manager_started.handle_acquired_tags = MagicMock() ingestor_manager_started.ingestor_manager.managed_tags = True ingestor_manager_started.manage_no_slots(0) ingestor_manager_started.ingestor_manager.get_slot_leases.assert_not_called() ingestor_manager_started.ingestor_manager.handle_acquired_tags.assert_not_called() def test_manage_slots_none_available(ingestor_manager_started): ingestor_manager_started.handle_acquired_tags = MagicMock() ingestor_manager_started.ingestor_manager.managed_tags = False ingestor_manager_started.manage_no_slots(0) ingestor_manager_started.ingestor_manager.get_slot_leases.assert_not_called() ingestor_manager_started.ingestor_manager.handle_acquired_tags.assert_not_called() @mark.asyncio async def test_manage_slots_none_available_none_available(ingestor_manager_started): ingestor_manager_started.handle_acquired_tags = MagicMock() ingestor_manager_started.ingestor_manager.managed_tags = False await ingestor_manager_started.manage_no_slots(2) ingestor_manager_started.ingestor_manager.get_slot_leases.assert_called_once_with(1) @mark.asyncio async def test_manage_leases_no_ingestor_manager(ingestor_manager_started): ingestor_manager_started.ingestor_manager = None assert await ingestor_manager_started.manage_leases(2, 2, 5) is None @mark.asyncio async def test_manage_leases_no_available_slots_no_extra_slots(ingestor_manager_started): ingestor_manager_started.handle_acquired_tags = MagicMock() await ingestor_manager_started.manage_leases(0, 0, 0) ingestor_manager_started.ingestor_manager.get_slot_leases.assert_not_called() ingestor_manager_started.ingestor_manager.handle_acquired_tags.assert_not_called() ingestor_manager_started.ingestor_manager.drop_slot_leases.assert_not_called() @mark.asyncio async def test_manage_leases_available_slots_innactive_ingestors(ingestor_manager_started): ingestor_manager_started.handle_acquired_tags = MagicMock() await ingestor_manager_started.manage_leases(2, 2, 5) ingestor_manager_started.ingestor_manager.get_slot_leases.assert_called_once_with(2) ingestor_manager_started.ingestor_manager.drop_slot_leases.assert_not_called() @mark.asyncio async def test_manage_leases_no_available_slots_extra_sltos(ingestor_manager_started): ingestor_manager_started.handle_acquired_tags = MagicMock() ingestor_manager_started.ingestor_manager.managed_tags = { 'tag1': 'server1', 'tag2': 'server2', 'tag3': 'server3', } await ingestor_manager_started.manage_leases(0, 0, 2) ingestor_manager_started.ingestor_manager.get_slot_leases.assert_not_called() ingestor_manager_started.handle_acquired_tags.assert_not_called() ingestor_manager_started.ingestor_manager.drop_slot_leases.assert_called_once_with( ['tag2', 'tag3'] ) @mark.asyncio async def test_loop(ingestor_manager_started): ingestor_manager_started.manage_no_slots = AsyncMock() ingestor_manager_started.manage_leases = AsyncMock() ingestor_manager_started.update_ingestor_manager = AsyncMock() ingestor_manager_started.ingestor_manager.check_opc_servers_integrity = AsyncMock() ingestor_manager_started.ingestor_manager.managed_tags = { 'slot1': 'server1', 'slot2': 'server2', 'slot3': 'server3', } ingestor_manager_started.ingestor_manager.get_active_ingestors = AsyncMock( return_value=['ingestor1', 'ingestor2'] ) ingestor_manager_started.ingestor_manager.get_number_of_slots = AsyncMock(return_value=5) ingestor_manager_started.ingestor_manager.get_number_of_leases = AsyncMock(return_value=1) await ingestor_manager_started.loop() ingestor_manager_started.ingestor_manager.declare_active.assert_called_once() ingestor_manager_started.ingestor_manager.get_active_ingestors.assert_called_once() ingestor_manager_started.ingestor_manager.get_number_of_slots.assert_called_once() ingestor_manager_started.manage_no_slots.assert_called_once_with( ingestor_manager_started.ingestor_manager.get_number_of_slots.return_value ) # Explanation: 5 - 2 = 3, 3 - 1 = 2 ingestor_manager_started.manage_leases.assert_called_once_with(4, 3, 2) ingestor_manager_started.ingestor_manager.update_slot_config.assert_called_once() ingestor_manager_started.ingestor_manager.check_opc_servers_integrity.assert_called_once() ingestor_manager_started.update_ingestor_manager.assert_called_once() @mark.asyncio async def test_loop_no_managed(ingestor_manager_started): ingestor_manager_started.manage_no_slots = AsyncMock() ingestor_manager_started.manage_leases = AsyncMock() ingestor_manager_started.update_ingestor_manager = AsyncMock() ingestor_manager_started.ingestor_manager.check_opc_servers_integrity = AsyncMock() ingestor_manager_started.ingestor_manager.managed_tags = {} ingestor_manager_started.ingestor_manager.get_active_ingestors = AsyncMock( return_value=['ingestor1', 'ingestor2'] ) ingestor_manager_started.ingestor_manager.get_number_of_slots = AsyncMock(return_value=5) await ingestor_manager_started.loop() ingestor_manager_started.ingestor_manager.declare_active.assert_called_once() ingestor_manager_started.ingestor_manager.get_active_ingestors.assert_called_once() ingestor_manager_started.ingestor_manager.get_number_of_slots.assert_called_once() ingestor_manager_started.manage_no_slots.assert_called_once_with( ingestor_manager_started.ingestor_manager.get_number_of_slots.return_value ) # Explanation: 5 - 2 = 3, 3 - 1 = 2 ingestor_manager_started.manage_leases.assert_called_once_with(ANY, 3, -1) ingestor_manager_started.ingestor_manager.update_slot_config.assert_called_once() ingestor_manager_started.logger.info.assert_any_call('No slots acquired in this loop') ingestor_manager_started.ingestor_manager.check_opc_servers_integrity.assert_called_once() ingestor_manager_started.update_ingestor_manager.assert_called_once() @mark.asyncio async def test_loop_no_ingestor_manager(ingestor_manager_started): ingestor_manager_started.ingestor_manager = None assert await ingestor_manager_started.loop() is None @mark.asyncio async def test_update_ingestor_manager(ingestor_manager_started): ingestor_manager_started.ingestor_manager.managed_tags = { 'slot_to_create': 'new_config', 'slot_to_update': 'new_config', 'slot_to_do_nothing': 'old_config', } old_managed_tags = { 'slot_to_update': 'old_config', 'slot_to_delete': 'old_config', 'slot_to_do_nothing': 'old_config', } await ingestor_manager_started.update_ingestor_manager(old_managed_tags) ingestor_manager_started.ingestor_manager.update_opc_servers.assert_called_once() ingestor_manager_started.ingestor_manager.subscribe_to_tags.assert_has_calls( [call({'slot_to_create': 'new_config'}), call({'slot_to_update': 'new_config'})] ) ingestor_manager_started.ingestor_manager.unsubscribe_slot.assert_has_calls( [call('slot_to_update'), call('slot_to_delete')] ) @mark.asyncio async def test_update_ingestor_manager_no_ingestor_manager(ingestor_manager_started): ingestor_manager_started.ingestor_manager = None await ingestor_manager_started.update_ingestor_manager({})