Files
sientia-dataops-opc-ingestor/tests/unit/test_ingestor.py
vitor-aignosi 6d2a0e46be Refactor OPC server integrity checks to be asynchronous
Updated the Ingestor and IngestorManager classes to use async methods for checking OPC server integrity and shutting down OPC managers. Adjusted related tests to ensure proper mocking of asynchronous behavior.
2025-10-23 11:46:02 -03:00

337 lines
13 KiB
Python

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
'localhost1', # REDIS_HOST
'63790', # REDIS_PORT
'user', # REDIS_USERNAME
'password', # REDIS_PASSWORD
'100', # LEASE_TTL
'200', # HEARTBEAT_TTL
'localhost1', # 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 == 'localhost1'
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 == 'localhost1'
assert ingestor.poll_interval == 50
assert ingestor.metadata == {
'model_id': '-',
'model_name': '-',
'workflow_name': 'opc_ingestor',
'schema_name': 'opc_ingestor',
'pod_id': 'localhost1',
}
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 = MagicMock(
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.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,
)
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()
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
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 = MagicMock()
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 = MagicMock(
return_value=['ingestor1', 'ingestor2']
)
ingestor_manager_started.ingestor_manager.get_number_of_slots = MagicMock(return_value=5)
ingestor_manager_started.ingestor_manager.get_number_of_leases = MagicMock(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 = MagicMock()
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 = MagicMock(
return_value=['ingestor1', 'ingestor2']
)
ingestor_manager_started.ingestor_manager.get_number_of_slots = MagicMock(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({})