sientia-dataops-opc-ingestor/tests/unit/test_ingestor.py

from copy import deepcopy
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch

from pytest import fixture, mark

from ingestor.ingestor import Ingestor
from ingestor.managers.ingestor_manager import IngestorManager


@patch('ingestor.ingestor.getenv')
@patch('ingestor.ingestor.NotificationHandler')
def test___init__(notification_handler, getenv):
    getenv.side_effect = [
        '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('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.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(
        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,
        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']
    )
    # The extra slots are unsubscribed before their leases are given back
    calls = [c[0] for c in ingestor_manager_started.ingestor_manager.mock_calls]
    assert calls.index('drop_slot_leases') > calls.index('unsubscribe_slot')
    assert list(ingestor_manager_started.ingestor_manager.managed_tags) == ['tag1']


@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({})


def _manager_with_real_unsubscribe(managed_tags, readers):
    """A mocked IngestorManager whose unsubscribe path (unsubscribe_slot and
    _unsubscribe_reader) is the real one."""
    manager = AsyncMock(update_opc_servers=AsyncMock(), subscribe_to_tags=AsyncMock())
    manager.managed_tags = managed_tags
    manager.opc_managers = readers
    manager.unsubscribe_slot = lambda slot: IngestorManager.unsubscribe_slot(manager, slot)
    manager._unsubscribe_reader = lambda server, reader, slot: IngestorManager._unsubscribe_reader(
        manager, server, reader, slot
    )
    return manager


@mark.asyncio
async def test_dropped_extra_lease_does_not_stop_the_loop(ingestor_manager_started, reader_double):
    """Regression: manage_leases drops slot 2, then update_ingestor_manager sees slot 2 in
    the snapshot only and unsubscribes it again. That used to raise KeyError and stop
    the main loop (container restart)."""
    reader = reader_double('1', '2')
    ingestor_manager_started.ingestor_manager = _manager_with_real_unsubscribe(
        {'1': {'kepserver': {}}, '2': {'kepserver': {}}}, {'kepserver': reader}
    )
    snapshot = deepcopy(ingestor_manager_started.ingestor_manager.managed_tags)

    await ingestor_manager_started.manage_leases(0, 0, 1)
    await ingestor_manager_started.update_ingestor_manager(snapshot)

    reader.unsubscribe.assert_called_once_with('2')
    assert list(reader.subscriptions) == ['1']
    assert list(ingestor_manager_started.ingestor_manager.managed_tags) == ['1']


@mark.asyncio
async def test_slot_removed_from_redis_is_unsubscribed(ingestor_manager_started, reader_double):
    """Regression: update_slot_config drops a slot deleted by the orchestrator from
    managed_tags without unsubscribing it; update_ingestor_manager must close its
    subscription instead of raising KeyError."""
    reader = reader_double('1', '3')
    ingestor_manager_started.ingestor_manager = _manager_with_real_unsubscribe(
        {'1': {'kepserver': {}}}, {'kepserver': reader}
    )
    snapshot = {'1': {'kepserver': {}}, '3': {'kepserver': {}}}

    await ingestor_manager_started.update_ingestor_manager(snapshot)

    reader.unsubscribe.assert_called_once_with('3')
    assert list(reader.subscriptions) == ['1']