sientia-dataops-opc-ingestor/tests/unit/managers/test_resource_manager.py

from unittest.mock import AsyncMock, MagicMock, patch

from pytest import fixture, mark, raises

from ingestor.managers.resource_manager import ResourceManager

metadata = {
    'metadata': {
        'model_id': 'test_model',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
        'schema_name': 'test_schedule',
    },
}


@patch('ingestor.managers.resource_manager.RedisRepository')
def test___init__(redis_repository):
    logger = MagicMock()
    notification_handler = MagicMock()
    metrics_controller = MagicMock()
    resource_manager = ResourceManager(
        host='localhost',
        port=6379,
        lease_ttl=10,
        heartbeat_ttl=10,
        metadata=metadata['metadata'],
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
    )
    assert resource_manager.redis_repository == redis_repository.return_value
    redis_repository.assert_called_once_with(
        host='localhost',
        port=6379,
        username=None,
        password=None,
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
    )
    redis_repository.return_value.redis_client.ping.assert_called_once()


@patch('ingestor.managers.resource_manager.RedisRepository')
def test___init__connection_failure(redis_repository):
    redis_repository.return_value.redis_client.ping.side_effect = Exception('Connection failed')
    with raises(Exception, match='Connection failed'):
        ResourceManager(
            host='localhost',
            port=6379,
            lease_ttl=10,
            heartbeat_ttl=10,
            metadata=metadata['metadata'],
            logger=MagicMock(),
            notification_handler=MagicMock(),
            metrics_controller=MagicMock(),
        )
        redis_repository.return_value.redis_client.ping.assert_called_once()
        redis_repository.logger.error.assert_called_once_with(
            'Failed to connect to Redis: Connection failed'
        )
        redis_repository.return_value.redis_client.ping.assert_called_once()


@fixture
@patch('ingestor.managers.resource_manager.RedisRepository')
def resource_manager(redis_repository):
    resource_manager = ResourceManager(
        host='localhost',
        port=6379,
        lease_ttl=10,
        heartbeat_ttl=10,
        metadata=metadata['metadata'],
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=MagicMock(),
    )

    resource_manager.send_notification = MagicMock()
    resource_manager.send_notification_async = AsyncMock()
    resource_manager.emit_metric = AsyncMock()

    resource_manager.redis_repository = AsyncMock()

    return resource_manager


@mark.asyncio
async def test_get_tag_slot(resource_manager):
    result = await resource_manager.get_tag_slot('id')
    assert result == resource_manager.redis_repository.get.return_value

    resource_manager.redis_repository.get.assert_called_once_with(
        'slot:opc_tags:id', metadata=metadata['metadata']
    )


@mark.asyncio
async def test_ingestor_heartbeat(resource_manager):
    await resource_manager.ingestor_heartbeat()
    resource_manager.redis_repository.set.assert_called_once_with(
        'heartbeat:ingestor:localhost', 1, ttl=10, metadata=metadata['metadata']
    )


@mark.asyncio
async def test_lease_tag(resource_manager):
    output = await resource_manager.lease_tag('tag_id')
    assert output is resource_manager.redis_repository.set.return_value
    resource_manager.redis_repository.set.assert_called_once_with(
        'lease:opc_tags:tag_id', 'localhost', ttl=10, nx=True, metadata=metadata['metadata']
    )


@mark.asyncio
async def test_renew_tag_lease_success(resource_manager):
    resource_manager.redis_repository.get.return_value = 'localhost'
    resource_manager.redis_repository.expire.return_value = True
    result = await resource_manager.renew_tag_lease('tag_id')
    assert result is resource_manager.redis_repository.expire.return_value
    resource_manager.redis_repository.get.assert_called_once_with(
        'lease:opc_tags:tag_id', metadata=metadata['metadata']
    )
    resource_manager.redis_repository.expire.assert_called_once_with(
        'lease:opc_tags:tag_id', 10, metadata=metadata['metadata']
    )


@mark.asyncio
async def test_renew_tag_lease_failure(resource_manager):
    resource_manager.redis_repository.get.return_value = 'other_pod_id'

    result = await resource_manager.renew_tag_lease('tag_id')
    assert result is False

    resource_manager.redis_repository.get.assert_called_once_with(
        'lease:opc_tags:tag_id', metadata=metadata['metadata']
    )
    resource_manager.redis_repository.expire.assert_not_called()


@mark.asyncio
async def test_drop_tag_lease(resource_manager):
    await resource_manager.drop_tag_lease('tag_id')
    resource_manager.redis_repository.delete.assert_called_once_with(
        'lease:opc_tags:tag_id', metadata=metadata['metadata']
    )


@mark.asyncio
async def test_get_all_ingestors(resource_manager):
    resource_manager.redis_repository.keys.return_value = ['ingestor1', 'ingestor2']
    result = await resource_manager.get_all_ingestors()
    assert result == ['ingestor1', 'ingestor2']
    resource_manager.redis_repository.keys.assert_called_once_with(
        'heartbeat:ingestor:*', metadata=metadata['metadata']
    )


@mark.asyncio
async def test_get_all_slots(resource_manager):
    resource_manager.redis_repository.keys.return_value = ['slot1', 'slot2']
    result = await resource_manager.get_all_slots()
    assert result == ['slot1', 'slot2']
    resource_manager.redis_repository.keys.assert_called_once_with(
        'slot:opc_tags:*', metadata=metadata['metadata']
    )


@mark.asyncio
async def test_get_all_leases(resource_manager):
    resource_manager.redis_repository.keys.return_value = ['lease1', 'lease2']
    result = await resource_manager.get_all_leases()
    assert result == ['lease1', 'lease2']
    resource_manager.redis_repository.keys.assert_called_once_with(
        'lease:opc_tags:*', metadata=metadata['metadata']
    )