90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from ingestor.engines.write import WriteResult
|
|
from ingestor.io.data_sink import DataSink
|
|
|
|
|
|
def make_sink():
|
|
postgres = MagicMock()
|
|
connection = MagicMock()
|
|
postgres.engine.begin.return_value.__enter__.return_value = connection
|
|
sink = DataSink(data_manager=MagicMock(), postgres=postgres)
|
|
return sink, connection
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_write_status_targets_model_row():
|
|
sink, connection = make_sink()
|
|
cmd = {
|
|
'schema': 'sientia_data',
|
|
'table_name': 'pred_m1',
|
|
'model_id': '1',
|
|
'timestamp': '2026-07-13T00:00:00Z',
|
|
}
|
|
result = WriteResult(ok=True, status='ok', opc_status_code='Good', message='')
|
|
|
|
await sink.update_write_status(cmd, result)
|
|
|
|
assert connection.execute.call_count == 1
|
|
stmt, params = connection.execute.call_args[0]
|
|
sql = str(stmt)
|
|
assert 'update' in sql.lower() and 'write_status' in sql.lower()
|
|
assert 'sientia_data.pred_m1' in sql
|
|
# values must be bound params, never interpolated into the SQL text itself
|
|
assert params == {'status': 'ok', 'timestamp': '2026-07-13T00:00:00Z', 'model_id': '1'}
|
|
assert 'ok' not in sql
|
|
assert '2026-07-13T00:00:00Z' not in sql
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_write_status_rejects_bad_schema():
|
|
sink, connection = make_sink()
|
|
cmd = {
|
|
'schema': 'a; DROP TABLE x',
|
|
'table_name': 'pred_m1',
|
|
'model_id': '1',
|
|
'timestamp': '2026-07-13T00:00:00Z',
|
|
}
|
|
result = WriteResult(ok=False, status='error', opc_status_code='Bad', message='x')
|
|
|
|
with pytest.raises(ValueError):
|
|
await sink.update_write_status(cmd, result)
|
|
|
|
connection.execute.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_write_status_rejects_trailing_newline_identifier():
|
|
sink, connection = make_sink()
|
|
cmd = {
|
|
'schema': 'sientia_data\n',
|
|
'table_name': 'pred_m1',
|
|
'model_id': '1',
|
|
'timestamp': '2026-07-13T00:00:00Z',
|
|
}
|
|
result = WriteResult(ok=True, status='ok', opc_status_code='Good', message='')
|
|
|
|
with pytest.raises(ValueError):
|
|
await sink.update_write_status(cmd, result)
|
|
|
|
connection.execute.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_write_status_rejects_bad_table_name():
|
|
sink, connection = make_sink()
|
|
cmd = {
|
|
'schema': 'sientia_data',
|
|
'table_name': 'pred_m1; DROP TABLE x',
|
|
'model_id': '1',
|
|
'timestamp': '2026-07-13T00:00:00Z',
|
|
}
|
|
result = WriteResult(ok=False, status='error', opc_status_code='Bad', message='x')
|
|
|
|
with pytest.raises(ValueError):
|
|
await sink.update_write_status(cmd, result)
|
|
|
|
connection.execute.assert_not_called()
|