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

SECURITY = {'mode': 'SignAndEncrypt', 'policy': 'Basic256Sha256'}
CERTIFICATE = {
    'certificate_id': 7,
    'application_uri': 'urn:sientia:opc-ingestor',
    'cert': {'bucket': 'opc-certificates', 'object_name': 'client/p/cert.pem'},
    'private_key': {'bucket': 'opc-certificates', 'object_name': 'client/p/key.pem'},
}
SECURED = {'url': 'opc.tcp://h:4840', 'security': SECURITY, 'certificate': CERTIFICATE}


def make_pool(cert_provider=None):
    pool = OpcConnectionPool(
        cert_provider=cert_provider or AsyncMock(),
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
    )
    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', 'certificate': None})
        c2 = await pool.get('srv', {'url': 'opc.tcp://h:4840', 'certificate': None})

    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', 'certificate': None})

    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', 'certificate': None})
    assert client is fake_client2
    assert fake_client2.connect.await_count == 1


@pytest.mark.asyncio
async def test_get_applies_the_certificate_and_security_of_the_server():
    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', SECURED)

    cert_provider.fetch.assert_awaited_once_with(CERTIFICATE)
    mock_apply.assert_awaited_once_with(fake_client, SECURITY, CERTIFICATE, 'cert.pem', 'key.pem')


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'cfg',
    [
        pytest.param({**SECURED, 'certificate': None}, id='no-certificate'),
        pytest.param({**SECURED, 'security': {'mode': 'None', 'policy': 'None'}}, id='mode-none'),
    ],
)
async def test_get_without_certificate_connects_without_security(cfg):
    """No client certificate (or mode None): no security, no MinIO download (D5)"""
    cert_provider = AsyncMock()
    pool = make_pool(cert_provider=cert_provider)
    with (
        patch('ingestor.connection.pool.Client', return_value=AsyncMock()),
        patch('ingestor.connection.pool.apply_security', new=AsyncMock()) as mock_apply,
    ):
        await pool.get('srv', cfg)

    cert_provider.fetch.assert_not_awaited()
    mock_apply.assert_not_awaited()


@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', 'certificate': None})

    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', 'certificate': None})

    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', 'certificate': None})
        await pool.get('b', {'url': 'opc.tcp://h:4840', 'certificate': None})

    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', 'certificate': None})

    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', 'certificate': None})

    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', 'certificate': None}
    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


@pytest.mark.asyncio
async def test_release_idle_keeps_pinned_connection():
    """A subscribed connection only gets callbacks (never get()), so it looks idle:
    the pin keeps release_idle from dropping it (SIENTIAPDE-1974 G4)."""
    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', 'certificate': None})

    pool.pin('srv')
    await pool.release_idle(max_idle_s=0)

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


@pytest.mark.asyncio
async def test_release_idle_evicts_after_last_unpin():
    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', 'certificate': None})

    pool.pin('srv')
    pool.pin('srv')  # two slots subscribed on the same server
    pool.unpin('srv')
    await pool.release_idle(max_idle_s=0)
    assert 'srv' in pool._connections  # one subscription still active

    pool.unpin('srv')
    await pool.release_idle(max_idle_s=0)
    assert 'srv' not in pool._connections
    fake_client.disconnect.assert_awaited_once()


def test_unpin_never_goes_negative():
    pool = make_pool()

    pool.unpin('srv')
    pool.pin('srv')

    assert pool._pins == {'srv': 1}


@pytest.mark.asyncio
async def test_evict_disconnects_now_and_next_get_reconnects():
    pool = make_pool()
    first, second = AsyncMock(), AsyncMock()
    with patch('ingestor.connection.pool.Client', side_effect=[first, second]):
        await pool.get('srv', {'url': 'opc.tcp://h:4840', 'certificate': None})
        await pool.evict('srv')
        client = await pool.get('srv', {'url': 'opc.tcp://h:4840', 'certificate': None})

    first.disconnect.assert_awaited_once()
    assert client is second


@pytest.mark.asyncio
async def test_evict_unknown_server_is_noop():
    pool = make_pool()

    await pool.evict('nope')

    pool.emit_metric.assert_not_awaited()


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

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


@pytest.mark.asyncio
async def test_get_with_different_cfg_reuses_client_and_warns():
    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', 'certificate': None})
        client = await pool.get('srv', {'url': 'opc.tcp://10.0.0.5:4840', 'certificate': None})

    assert client is fake_client
    assert fake_client.connect.await_count == 1
    pool.logger.warning.assert_called_once()
    assert 'url' in pool.logger.warning.call_args.args[0]


@pytest.mark.asyncio
async def test_get_with_another_certificate_warns():
    cert_provider = AsyncMock()
    cert_provider.fetch.return_value = ('cert.pem', 'key.pem')
    pool = make_pool(cert_provider=cert_provider)
    with (
        patch('ingestor.connection.pool.Client', return_value=AsyncMock()),
        patch('ingestor.connection.pool.apply_security', new=AsyncMock()),
    ):
        await pool.get('srv', SECURED)
        await pool.get('srv', {**SECURED, 'certificate': {**CERTIFICATE, 'certificate_id': 8}})

    assert 'certificate' in pool.logger.warning.call_args.args[0]


@pytest.mark.asyncio
async def test_secured_connect_failure_invalidates_cert_cache():
    from unittest.mock import MagicMock

    cert_provider = MagicMock()
    cert_provider.fetch = AsyncMock(return_value=('cert.pem', 'key.pem'))
    pool = make_pool(cert_provider=cert_provider)
    fake_client = AsyncMock()
    fake_client.connect.side_effect = RuntimeError('BadCertificateUntrusted')
    with (
        patch('ingestor.connection.pool.Client', return_value=fake_client),
        patch('ingestor.connection.pool.apply_security', AsyncMock()),
    ):
        with pytest.raises(RuntimeError):
            await pool.get('srv', SECURED)

    cert_provider.invalidate.assert_called_once_with(7)


@pytest.mark.asyncio
async def test_insecure_connect_failure_keeps_cert_cache():
    from unittest.mock import MagicMock

    cert_provider = MagicMock()
    pool = make_pool(cert_provider=cert_provider)
    fake_client = AsyncMock()
    fake_client.connect.side_effect = RuntimeError('unreachable')
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        with pytest.raises(RuntimeError):
            await pool.get('srv', {'url': 'opc.tcp://h:4840', 'certificate': None})

    cert_provider.invalidate.assert_not_called()


@pytest.mark.asyncio
async def test_connect_identifies_the_pod_on_the_server():
    """Session/application name = pod id, as the legacy OpcManager did, so each pod's
    session is recognizable in the KEPServer console."""
    pool = make_pool()
    fake_client = MagicMock(connect=AsyncMock())
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        await pool.get('srv', {'url': 'opc.tcp://h:4840', 'certificate': None})

    assert fake_client.name == pool.pod_id
    assert fake_client.description == pool.pod_id
    assert fake_client.product_uri == pool.pod_id.replace('-', ':')
    assert fake_client.application_uri == pool.pod_id.replace('-', ':')


@pytest.mark.asyncio
async def test_connect_with_certificate_uses_its_application_uri():
    """apply_security sets the certificate's URI: it must match its SAN URI
    (BadCertificateUriInvalid otherwise)"""
    cert_provider = AsyncMock()
    cert_provider.fetch.return_value = ('cert.pem', 'key.pem')
    pool = make_pool(cert_provider=cert_provider)
    fake_client = MagicMock(connect=AsyncMock(), set_security=AsyncMock())
    with patch('ingestor.connection.pool.Client', return_value=fake_client):
        await pool.get('srv', SECURED)

    assert fake_client.application_uri == 'urn:sientia:opc-ingestor'
    assert fake_client.name == pool.pod_id


@pytest.mark.asyncio
async def test_config_for_returns_a_copy_of_the_live_connection_settings():
    pool = make_pool()
    cfg = {'url': 'opc.tcp://h:4840', 'security': SECURITY, 'certificate': None}
    with patch('ingestor.connection.pool.Client', return_value=AsyncMock()):
        await pool.get(1, cfg)

    found = pool.config_for('1')
    assert found == cfg
    found['url'] = 'mutated'
    assert pool.config_for(1)['url'] == 'opc.tcp://h:4840'

    await pool.evict('1')
    assert pool.config_for('1') is None


def test_config_for_unknown_server_is_none():
    assert make_pool().config_for('missing') is None