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