sientia-dataops-opc-ingestor/tests/unit/connection/test_pool.py

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

import ingestor.metrics as metrics
from ingestor.connection.pool import OpcConnectionPool


def make_pool(cert_provider=None):
    pool = OpcConnectionPool(
        cert_provider=cert_provider or AsyncMock(),
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
        metadata={},
    )
    pool.emit_metric = AsyncMock()
    return pool


@pytest.mark.asyncio
async def test_get_reuses_connection_per_server():
    pool = make_pool()
    fake_client = AsyncMock()
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        c1 = await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False})
        c2 = await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False})

    assert c1 is c2
    assert fake_client.connect.await_count == 1


@pytest.mark.asyncio
async def test_get_failure_records_metric_and_raises():
    pool = make_pool()
    fake_client = AsyncMock()
    fake_client.connect.side_effect = RuntimeError('boom')
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        with pytest.raises(RuntimeError):
            await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False})

    pool.emit_metric.assert_any_await(
        metric_object=metrics.OPC_CONNECTIONS_FAILED,
        tags={'pod_id': pool.pod_id, 'server_name': 'srv'},
    )
    assert 'srv' not in pool._connections

    # a later successful get must reconnect (failed attempt was not cached)
    fake_client2 = AsyncMock()
    with patch('ingestor.connection.pool.Client', return_value=fake_client2):
        client = await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False})
    assert client is fake_client2
    assert fake_client2.connect.await_count == 1


@pytest.mark.asyncio
async def test_get_applies_security_when_requested():
    cert_provider = AsyncMock()
    cert_provider.fetch.return_value = ('cert.pem', 'key.pem')
    pool = make_pool(cert_provider=cert_provider)
    fake_client = AsyncMock()
    with (
        patch('ingestor.connection.pool.Client', return_value=fake_client),
        patch('ingestor.connection.pool.apply_security', new=AsyncMock()) as mock_apply,
    ):
        await pool.get(
            'srv',
            {'url': 'opc.tcp://h:4840', 'server_uri': 'urn:test', 'use_security': True},
        )

    cert_provider.fetch.assert_awaited_once()
    mock_apply.assert_awaited_once_with(fake_client, 'urn:test', 'cert.pem', 'key.pem', None)


@pytest.mark.asyncio
async def test_release_idle_evicts_and_disconnects():
    pool = make_pool()
    fake_client = AsyncMock()
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False})

    client, _ = pool._connections['srv']
    pool._connections['srv'] = (client, -1000.0)  # force idle

    await pool.release_idle(max_idle_s=0)

    fake_client.disconnect.assert_awaited_once()
    assert 'srv' not in pool._connections
    pool.emit_metric.assert_any_await(
        metric_object=metrics.OPC_DISCONNECTIONS_TOTAL,
        tags={'pod_id': pool.pod_id, 'server_name': 'srv'},
    )


@pytest.mark.asyncio
async def test_release_idle_keeps_recently_used():
    pool = make_pool()
    fake_client = AsyncMock()
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False})

    await pool.release_idle(max_idle_s=3600)

    fake_client.disconnect.assert_not_awaited()
    assert 'srv' in pool._connections


@pytest.mark.asyncio
async def test_close_all_disconnects_everything():
    pool = make_pool()
    fake_a, fake_b = AsyncMock(), AsyncMock()
    with patch('ingestor.connection.pool.Client', side_effect=[fake_a, fake_b]):
        await pool.get('a', {'url': 'opc.tcp://h:4840', 'use_security': False})
        await pool.get('b', {'url': 'opc.tcp://h:4840', 'use_security': False})

    await pool.close_all()

    fake_a.disconnect.assert_awaited_once()
    fake_b.disconnect.assert_awaited_once()
    assert pool._connections == {}


@pytest.mark.asyncio
async def test_release_idle_disconnect_retries_then_gives_up():
    pool = make_pool()
    fake_client = AsyncMock()
    fake_client.disconnect.side_effect = RuntimeError('nope')
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False})

    client, _ = pool._connections['srv']
    pool._connections['srv'] = (client, -1000.0)

    await pool.release_idle(max_idle_s=0)

    assert fake_client.disconnect.await_count == 5
    assert 'srv' not in pool._connections


@pytest.mark.asyncio
async def test_success_path_emits_connection_metrics():
    pool = make_pool()
    fake_client = AsyncMock()
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False})

    pool.emit_metric.assert_any_await(
        metric_object=metrics.OPC_CONNECTIONS_TOTAL,
        tags={'pod_id': pool.pod_id, 'server_name': 'srv'},
    )
    pool.emit_metric.assert_any_await(
        metric_object=metrics.OPC_CONNECTION_STATUS,
        method='set',
        value=1,
        tags={'pod_id': pool.pod_id, 'server_name': 'srv', 'server_url': 'opc.tcp://h:4840'},
    )


@pytest.mark.asyncio
async def test_concurrent_get_same_server_connects_once():
    pool = make_pool()
    started = asyncio.Event()
    proceed = asyncio.Event()

    clients = [AsyncMock(), AsyncMock()]

    async def slow_connect():
        started.set()
        await proceed.wait()

    clients[0].connect.side_effect = slow_connect

    cfg = {'url': 'opc.tcp://h:4840', 'use_security': False}
    with patch('ingestor.connection.pool.Client', side_effect=clients) as ctor:
        t1 = asyncio.create_task(pool.get('srv', cfg))
        await started.wait()  # t1 is mid-connect, holding the per-server lock
        t2 = asyncio.create_task(pool.get('srv', cfg))
        await asyncio.sleep(0)  # let t2 block on the lock (not build a 2nd Client)
        proceed.set()
        c1, c2 = await asyncio.gather(t1, t2)

    assert c1 is c2 is clients[0]
    assert ctor.call_count == 1
    assert clients[0].connect.await_count == 1