sientia-dataops-opc-ingestor/tests/unit/io/test_data_sink_status.py

from unittest.mock import AsyncMock, MagicMock

import pytest
from sqlalchemy import create_engine, text
from sqlalchemy.dialects import postgresql

import ingestor.metrics as metrics
from ingestor.engines.write import WriteResult
from ingestor.io.data_sink import DataSink, write_status_statements

CMD = {
    'server_id': '1',
    'addr': 'ns=2;s=Pred',
    'schema': 'sientia_data',
    'table_name': 'pred_m1',
    'model_id': '1',
    'timestamp': '2026-07-13T00:00:00Z',
}


def _result(status):
    return WriteResult(ok=status == 'ok', status=status, opc_status_code='x', message='')


def make_sink(rowcount=1, existing_row=True):
    postgres = MagicMock()
    connection = MagicMock()
    connection.execute.return_value.rowcount = rowcount
    connection.execute.return_value.first.return_value = (1,) if existing_row else None
    postgres.engine.begin.return_value.__enter__.return_value = connection
    data_manager = MagicMock()
    data_manager.emit_metric = AsyncMock()
    sink = DataSink(data_manager=data_manager, postgres=postgres)
    return sink, connection


def _outcome_metric(sink):
    calls = [
        c
        for c in sink.data_manager.emit_metric.await_args_list
        if c.kwargs['metric_object'] is metrics.OPC_WRITE_STATUS_UPDATE_TOTAL
    ]
    assert len(calls) == 1
    return calls[0].kwargs['tags']['result']


@pytest.mark.asyncio
async def test_update_write_status_targets_model_row_with_precedence():
    sink, connection = make_sink()

    outcome = await sink.update_write_status(CMD, _result('ok'))

    assert outcome == 'updated'
    assert connection.execute.call_count == 1
    (stmt,) = connection.execute.call_args[0]
    compiled = stmt.compile(dialect=postgresql.dialect())
    sql = str(compiled)
    assert sql.startswith('UPDATE sientia_data.pred_m1 SET write_status=')
    assert 'CASE sientia_data.pred_m1.write_status WHEN' in sql
    # values must be bound params, never interpolated into the SQL text itself
    assert '2026-07-13T00:00:00Z' not in sql and "'ok'" not in sql
    assert compiled.params['write_status'] == 'ok'
    assert '2026-07-13T00:00:00Z' in compiled.params.values()
    assert _outcome_metric(sink) == 'updated'


@pytest.mark.asyncio
async def test_update_write_status_kept_when_the_row_has_a_worse_outcome():
    sink, _ = make_sink(rowcount=0, existing_row=True)

    assert await sink.update_write_status(CMD, _result('ok')) == 'kept'
    assert _outcome_metric(sink) == 'kept'


@pytest.mark.asyncio
async def test_update_write_status_row_not_found_is_logged():
    """E.g. the Laborious INSERT failed and it wrote to OPC anyway (SIENTIAPDE-1975 D10)."""
    sink, _ = make_sink(rowcount=0, existing_row=False)

    assert await sink.update_write_status(CMD, _result('error')) == 'row_not_found'
    assert 'no prediction row' in sink.data_manager.logger.warning.call_args.args[0]
    assert _outcome_metric(sink) == 'row_not_found'


@pytest.mark.asyncio
async def test_update_write_status_database_error_is_not_raised():
    sink, connection = make_sink()
    connection.execute.side_effect = RuntimeError('column "write_status" does not exist')

    assert await sink.update_write_status(CMD, _result('ok')) == 'error'
    assert 'write_status' in sink.data_manager.logger.warning.call_args.args[0]
    assert _outcome_metric(sink) == 'error'


@pytest.mark.parametrize(
    ('field', 'value'),
    [
        ('schema', 'a; DROP TABLE x'),
        ('schema', 'sientia_data\n'),
        ('table_name', 'pred_m1; DROP TABLE x'),
        ('schema', None),
        ('timestamp', None),
        ('model_id', None),
    ],
)
@pytest.mark.asyncio
async def test_update_write_status_rejects_invalid_commands(field, value):
    sink, connection = make_sink()

    outcome = await sink.update_write_status({**CMD, field: value}, _result('error'))

    assert outcome == 'invalid'
    connection.execute.assert_not_called()
    assert _outcome_metric(sink) == 'invalid'


def _apply(initial, answers):
    """Run the real statements on SQLite for each answer; return the final status."""
    engine = create_engine('sqlite://')
    with engine.begin() as connection:
        connection.execute(
            text('CREATE TABLE pred (model_id INTEGER, timestamp TEXT, write_status TEXT)')
        )
        connection.execute(
            text('INSERT INTO pred VALUES (1, :ts, :status)'),
            {'ts': '2026-07-13T00:00:00Z', 'status': initial},
        )
        for answer in answers:
            update_stmt, exists_stmt = write_status_statements(
                None, 'pred', answer, '2026-07-13T00:00:00Z', 1
            )
            connection.execute(update_stmt)
            assert connection.execute(exists_stmt).first() is not None
        return connection.execute(text('SELECT write_status FROM pred')).scalar()


@pytest.mark.parametrize(
    ('initial', 'answers', 'expected'),
    [
        ('pending', ['ok'], 'ok'),
        ('pending', ['error', 'ok'], 'error'),
        ('pending', ['ok', 'error'], 'error'),
        ('pending', ['session_bad', 'error', 'ok'], 'session_bad'),
        ('error', ['ok'], 'error'),
        (None, ['ok'], 'ok'),
    ],
)
def test_update_statement_keeps_the_worst_answer(initial, answers, expected):
    assert _apply(initial, answers) == expected