sientia-dataops-opc-ingestor/tests/unit/adapters/test_nats_consumer.py

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

import pytest

from ingestor.adapters.nats_consumer import NatsWriteConsumer


def make_consumer(**kwargs):
    return NatsWriteConsumer(
        nats_url='nats://localhost:4222',
        write_engine=AsyncMock(),
        logger=MagicMock(),
        **kwargs,
    )


@pytest.mark.asyncio
async def test_handle_valid_json_dispatches_to_write_engine():
    consumer = make_consumer()
    payload = {'server_id': 'srv1', 'addr': 'ns=2;i=5', 'value': 1.0, 'data_type': 'float'}
    msg = MagicMock(data=json.dumps(payload).encode())

    await consumer._handle(msg)

    consumer.write_engine.write.assert_awaited_once_with(payload)


@pytest.mark.asyncio
async def test_handle_invalid_json_drops_message_and_logs_error():
    consumer = make_consumer()
    msg = MagicMock(data=b'not-json')

    await consumer._handle(msg)

    consumer.write_engine.write.assert_not_awaited()
    consumer.logger.error.assert_called_once()


CLIENT = 'ingestor.adapters.nats_consumer.Client'


def _fake_nc(connect=None):
    nc = MagicMock()
    nc.connect = connect or AsyncMock()
    nc.subscribe = AsyncMock()
    nc.close = AsyncMock()
    nc.drain = AsyncMock()
    nc.is_closed = False
    return nc


async def _hang(*_args, **_kwargs):
    await asyncio.sleep(10)


@pytest.mark.asyncio
async def test_start_subscribes_with_subject_and_queue():
    consumer = make_consumer(subject='opc.write', queue='gateway-writers')
    nc = _fake_nc()

    with patch(CLIENT, return_value=nc):
        await consumer.start()

    kwargs = nc.connect.await_args.kwargs
    assert kwargs['servers'] == 'nats://localhost:4222'
    # Reconnection is unbounded once connected: an outage never closes it for good.
    assert kwargs['max_reconnect_attempts'] == -1
    nc.subscribe.assert_awaited_once_with('opc.write', queue='gateway-writers', cb=consumer._handle)
    assert consumer.closed is False


@pytest.mark.asyncio
async def test_start_gives_up_after_the_connect_timeout():
    """start() runs on the main loop: an unreachable NATS must not block it."""
    consumer = make_consumer(connect_timeout_s=0.1)
    nc = _fake_nc(connect=AsyncMock(side_effect=_hang))

    with patch(CLIENT, return_value=nc), pytest.raises(TimeoutError):
        await asyncio.wait_for(consumer.start(), timeout=5)

    nc.close.assert_awaited_once()
    assert consumer.closed is True


@pytest.mark.asyncio
async def test_start_closes_the_connection_when_subscribe_fails():
    consumer = make_consumer()
    nc = _fake_nc()
    nc.subscribe.side_effect = RuntimeError('permissions violation')

    with patch(CLIENT, return_value=nc), pytest.raises(RuntimeError):
        await consumer.start()

    nc.close.assert_awaited_once()
    assert consumer.closed is True


@pytest.mark.asyncio
async def test_handle_unexpected_engine_error_is_logged_not_raised():
    consumer = make_consumer()
    consumer.write_engine.write.side_effect = RuntimeError('bug')
    msg = MagicMock(data=json.dumps({'server_id': '1'}).encode())

    await consumer._handle(msg)

    assert 'bug' in consumer.logger.error.call_args.args[0]


@pytest.mark.asyncio
async def test_connection_events_are_logged_in_one_line():
    consumer = make_consumer()
    nc = _fake_nc()
    with patch(CLIENT, return_value=nc):
        await consumer.start()
    callbacks = nc.connect.await_args.kwargs

    await callbacks['error_cb'](OSError('connection refused'))
    await callbacks['disconnected_cb']()
    await callbacks['reconnected_cb']()

    warnings = [c.args[0] for c in consumer.logger.warning.call_args_list]
    assert warnings == [
        'NATS error: OSError connection refused',
        'Disconnected from NATS; reconnecting',
    ]


@pytest.mark.asyncio
async def test_stop_before_start_does_not_raise():
    consumer = make_consumer()

    await consumer.stop()  # no raise


@pytest.mark.asyncio
async def test_stop_after_start_drains_and_unsubscribes():
    consumer = make_consumer()
    nc = _fake_nc()
    sub = AsyncMock()
    nc.subscribe = AsyncMock(return_value=sub)

    with patch(CLIENT, return_value=nc):
        await consumer.start()
        await consumer.stop()

    sub.unsubscribe.assert_awaited_once()
    nc.drain.assert_awaited_once()


@pytest.mark.asyncio
async def test_stop_on_a_connection_closed_for_good_does_not_raise():
    consumer = make_consumer()
    nc = _fake_nc()
    with patch(CLIENT, return_value=nc):
        await consumer.start()
    nc.is_closed = True
    consumer.sub.unsubscribe = AsyncMock(side_effect=RuntimeError('connection closed'))

    await consumer.stop()  # no raise

    nc.drain.assert_not_awaited()
    assert consumer.closed is True