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

SERVER_CFG = {'url': 'opc.tcp://kepserver:49320', 'server_uri': 'urn:kep', 'use_security': True}


def make_resolver(cfg=SERVER_CFG):
    resolver = MagicMock()
    resolver.resolve = AsyncMock(return_value=cfg)
    return resolver


def make_engine(**kwargs):
    kwargs.setdefault('servers', make_resolver())
    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()


# --- SIENTIAPDE-1975: server resolution (GW1) and error reporting (GW4) ---


@pytest.mark.asyncio
async def test_write_uses_the_resolved_server_connection():
    """The C2 command of the Laborious carries only the server_id."""
    eng = make_engine()
    eng.pool.get.return_value = make_client()

    await eng.write(base_cmd())

    eng.servers.resolve.assert_awaited_once_with('srv1', base_cmd())
    eng.pool.get.assert_awaited_once_with('srv1', SERVER_CFG)


@pytest.mark.asyncio
async def test_write_unknown_server_records_error_without_connecting():
    eng = make_engine(servers=make_resolver(cfg=None))

    result = await eng.write(base_cmd())

    assert result.status == 'error'
    assert result.opc_status_code == 'server_not_found'
    eng.pool.get.assert_not_awaited()
    eng.sink.update_write_status.assert_awaited_once_with(base_cmd(), result)


@pytest.mark.asyncio
async def test_write_server_lookup_failure_records_error():
    resolver = make_resolver()
    resolver.resolve.side_effect = RuntimeError('mongo down')
    eng = make_engine(servers=resolver)

    result = await eng.write(base_cmd())

    assert result.status == 'error'
    assert result.opc_status_code == 'server_lookup'
    assert 'mongo down' in result.message
    eng.pool.get.assert_not_awaited()


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

    await eng.write(base_cmd(url='opc.tcp://h:4840', server_uri='urn:h', use_security=True))

    eng.pool.get.assert_awaited_once_with(
        'srv1', {'url': 'opc.tcp://h:4840', 'use_security': True, 'server_uri': 'urn:h'}
    )


@pytest.mark.asyncio
async def test_write_without_resolver_or_url_is_server_not_found():
    eng = make_engine(servers=None)

    result = await eng.write(base_cmd())

    assert result.opc_status_code == 'server_not_found'
    eng.pool.get.assert_not_awaited()


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

    await eng.write(base_cmd())

    logged = eng.logger.error.call_args.args[0]
    assert 'RuntimeError' in logged and 'BadWriteNotSupported' in logged
    assert 'addr=ns=2;i=5' in logged and 'model_id=m1' in logged
    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
    )
    # Same series as a success, so it shows the last outcome.
    assert status_call.kwargs['value'] == 0
    assert status_call.kwargs['tags'] == {'pod_id': eng.pod_id, 'model_id': 'm1', 'status': 'ok'}


@pytest.mark.asyncio
@pytest.mark.parametrize('timestamp', [None, 'not-a-timestamp'])
async def test_write_bad_timestamp_records_error_without_raising(timestamp):
    eng = make_engine(write_timestamp=True)
    cmd = base_cmd(timestamp=timestamp)
    if timestamp is None:
        del cmd['timestamp']

    result = await eng.write(cmd)

    assert result.status == 'error'
    assert result.opc_status_code == 'bad_timestamp'
    eng.pool.get.assert_not_awaited()
    eng.sink.update_write_status.assert_awaited_once()