sientia-dataops-opc-ingestor/tests/unit/engines/test_write.py

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

import pytest
from asyncua import ua

import ingestor.metrics as metrics
from ingestor.engines.write import WriteEngine


def make_engine(**kwargs):
    eng = WriteEngine(
        pool=AsyncMock(),
        sink=AsyncMock(),
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
        **kwargs,
    )
    eng.emit_metric = AsyncMock()
    return eng


def make_client(node=None):
    node = node or make_node()
    client = MagicMock()
    client.get_node = MagicMock(return_value=node)
    return client


def make_node():
    node = MagicMock()
    node.write_value = AsyncMock()
    return node


def base_cmd(**overrides):
    cmd = {
        'server_id': 'srv1',
        'addr': 'ns=2;i=5',
        'value': 42.0,
        'data_type': 'float',
        'tag_type': 'prediction',
        'model_id': 'm1',
        'model_name': 'model-1',
        'timestamp': '2026-07-13T12:00:00Z',
        'schema': 'sientia_data',
        'table_name': 'predictions_m1',
    }
    cmd.update(overrides)
    return cmd


@pytest.mark.asyncio
async def test_write_success_writes_datavalue_and_returns_ok():
    eng = make_engine()
    node = make_node()
    client = make_client(node)
    eng.pool.get.return_value = client

    result = await eng.write(base_cmd())

    client.get_node.assert_called_once_with('ns=2;i=5')
    node.write_value.assert_awaited_once()
    dv = node.write_value.await_args.args[0]
    assert isinstance(dv, ua.DataValue)
    assert dv.Value.Value == 42.0
    assert dv.Value.VariantType == ua.VariantType.Float
    assert result.ok is True
    assert result.status == 'ok'
    assert result.opc_status_code == 'Good'
    eng.sink.update_write_status.assert_awaited_once_with(base_cmd(), result)


@pytest.mark.asyncio
async def test_write_success_emits_metrics():
    eng = make_engine()
    eng.pool.get.return_value = make_client()

    await eng.write(base_cmd())

    eng.emit_metric.assert_any_await(
        metric_object=metrics.OPC_WRITE_TOTAL,
        tags={'pod_id': eng.pod_id, 'server_name': 'srv1', 'model_id': 'm1', 'result': 'success'},
    )
    status_call = next(
        c
        for c in eng.emit_metric.await_args_list
        if c.kwargs['metric_object'] is metrics.OPC_WRITE_STATUS_BY_MODEL
    )
    assert status_call.kwargs['value'] == 1
    assert status_call.kwargs['tags'] == {'pod_id': eng.pod_id, 'model_id': 'm1', 'status': 'ok'}
    latency_call = next(
        c
        for c in eng.emit_metric.await_args_list
        if c.kwargs['metric_object'] is metrics.OPC_WRITE_LATENCY
    )
    assert latency_call.kwargs['method'] == 'observe'
    assert latency_call.kwargs['value'] >= 0


@pytest.mark.asyncio
async def test_write_defaults_carry_statuscode_and_timestamp():
    eng = make_engine(write_timestamp=True, write_quality=True)
    node = make_node()
    eng.pool.get.return_value = make_client(node)

    await eng.write(base_cmd())

    dv = node.write_value.await_args.args[0]
    assert dv.StatusCode_ is not None
    assert dv.StatusCode_.is_good()
    assert dv.SourceTimestamp == datetime.fromisoformat('2026-07-13T12:00:00Z')


@pytest.mark.asyncio
async def test_write_engine_flags_off_omit_statuscode_and_timestamp():
    eng = make_engine(write_timestamp=False, write_quality=False)
    node = make_node()
    eng.pool.get.return_value = make_client(node)

    await eng.write(base_cmd())

    dv = node.write_value.await_args.args[0]
    assert dv.StatusCode_ is None
    assert dv.SourceTimestamp is None


@pytest.mark.asyncio
async def test_per_message_flag_override_disables_timestamp():
    # Engine default is True, but the command explicitly turns it off.
    eng = make_engine(write_timestamp=True, write_quality=True)
    node = make_node()
    eng.pool.get.return_value = make_client(node)

    await eng.write(base_cmd(write_timestamp=False))

    dv = node.write_value.await_args.args[0]
    assert dv.SourceTimestamp is None
    assert dv.StatusCode_ is not None  # write_quality untouched, stays on


@pytest.mark.asyncio
async def test_write_bad_type_returns_error_without_raising():
    eng = make_engine()

    result = await eng.write(base_cmd(value='abc', data_type='int'))

    assert result.ok is False
    assert result.status == 'error'
    eng.sink.update_write_status.assert_awaited_once()
    eng.pool.get.assert_not_awaited()
    eng.emit_metric.assert_any_await(
        metric_object=metrics.OPC_ERROR_BY_TYPE,
        tags={'pod_id': eng.pod_id, 'server_name': 'srv1', 'error_type': 'conversion'},
    )
    eng.emit_metric.assert_any_await(
        metric_object=metrics.OPC_WRITE_TOTAL,
        tags={'pod_id': eng.pod_id, 'server_name': 'srv1', 'model_id': 'm1', 'result': 'error'},
    )


@pytest.mark.asyncio
async def test_write_none_data_type_returns_error_without_raising():
    eng = make_engine()

    result = await eng.write(base_cmd(data_type=None))

    assert result.ok is False
    assert result.status == 'error'
    eng.sink.update_write_status.assert_awaited_once()
    eng.pool.get.assert_not_awaited()


@pytest.mark.asyncio
async def test_write_missing_value_key_returns_error_without_raising():
    eng = make_engine()
    cmd = base_cmd()
    del cmd['value']

    result = await eng.write(cmd)

    assert result.ok is False
    assert result.status == 'error'
    eng.sink.update_write_status.assert_awaited_once()
    eng.pool.get.assert_not_awaited()


@pytest.mark.asyncio
async def test_write_missing_data_type_key_returns_error_without_raising():
    eng = make_engine()
    cmd = base_cmd()
    del cmd['data_type']

    result = await eng.write(cmd)

    assert result.ok is False
    assert result.status == 'error'
    eng.sink.update_write_status.assert_awaited_once()
    eng.pool.get.assert_not_awaited()


@pytest.mark.asyncio
async def test_write_value_raises_maps_to_error_result_not_raised():
    eng = make_engine()
    node = make_node()
    node.write_value.side_effect = RuntimeError('boom')
    eng.pool.get.return_value = make_client(node)

    result = await eng.write(base_cmd())

    assert result.ok is False
    assert result.status == 'error'
    assert result.opc_status_code == 'RuntimeError'
    eng.sink.update_write_status.assert_awaited_once()
    eng.emit_metric.assert_any_await(
        metric_object=metrics.OPC_ERROR_BY_TYPE,
        tags={'pod_id': eng.pod_id, 'server_name': 'srv1', 'error_type': 'RuntimeError'},
    )


@pytest.mark.asyncio
async def test_write_session_bad_status_code_error_maps_to_session_bad():
    from asyncua.ua.uaerrors import BadSessionClosed

    eng = make_engine()
    node = make_node()
    node.write_value.side_effect = BadSessionClosed('closed')
    eng.pool.get.return_value = make_client(node)

    result = await eng.write(base_cmd())

    assert result.ok is False
    assert result.status == 'session_bad'


@pytest.mark.asyncio
async def test_rate_limit_sleeps_between_back_to_back_writes():
    eng = make_engine(min_pause_s=0.05)
    eng.pool.get.return_value = make_client()

    with patch('ingestor.engines.write.asyncio.sleep', new_callable=AsyncMock) as sleep_mock:
        await eng.write(base_cmd())
        sleep_mock.assert_not_awaited()
        await eng.write(base_cmd())
        sleep_mock.assert_awaited_once()
        remaining = sleep_mock.await_args.args[0]
        assert 0 < remaining <= 0.05


@pytest.mark.asyncio
async def test_no_rate_limit_when_min_pause_is_zero():
    eng = make_engine(min_pause_s=0.0)
    eng.pool.get.return_value = make_client()

    with patch('ingestor.engines.write.asyncio.sleep', new_callable=AsyncMock) as sleep_mock:
        await eng.write(base_cmd())
        await eng.write(base_cmd())
        sleep_mock.assert_not_awaited()