"""A recording stand-in for the import log's SQLAlchemy engine. The import's status writes are the part of the feature with the strictest contract: which table they name, which predicate they carry, which columns they touch, how many transactions they open and which of them may fail without failing the import. Asserting that against a mock whose `execute` returns a `MagicMock` proves nothing — the assertions need the statement text and the parameters. So this double records every statement it is handed, answers reads from a queue the test fills, and lets a test make a write fail on a chosen attempt. It deliberately understands nothing about SQL beyond the first word of the statement. """ from __future__ import annotations from typing import Any class RecordedStatement: """One statement as the engine saw it: its normalised text and its bound parameters.""" def __init__(self, statement: Any, parameters: dict[str, Any] | None) -> None: self.sql = ' '.join(str(statement).split()) self.parameters: dict[str, Any] = dict(parameters or {}) @property def verb(self) -> str: return self.sql.split(' ', 1)[0].upper() @property def is_write(self) -> bool: return self.verb in ('UPDATE', 'INSERT', 'DELETE') @property def table(self) -> str: """The relation the statement names, for the "one table only" assertions.""" words = self.sql.split() if self.verb == 'UPDATE': return words[1] for index, word in enumerate(words): if word.upper() in ('FROM', 'INTO') and index + 1 < len(words): return words[index + 1] return '' @property def assignments(self) -> set[str]: """The column names a `SET` clause assigns, so a test can assert what is never written.""" upper = self.sql.upper() if ' SET ' not in upper: return set() body = self.sql[upper.index(' SET ') + 5 :] where = body.upper().find(' WHERE ') if where != -1: body = body[:where] return {part.split('=')[0].strip() for part in body.split(',') if '=' in part} def __repr__(self) -> str: # pragma: no cover - debugging aid return f'RecordedStatement({self.sql!r}, {self.parameters!r})' class FakeResult: """What `Connection.execute` returns: a rowcount for writes, mappings for reads.""" def __init__(self, rowcount: int = 1, rows: list[dict[str, Any]] | None = None) -> None: self.rowcount = rowcount self._rows = rows or [] def mappings(self) -> FakeResult: return self def first(self) -> dict[str, Any] | None: return self._rows[0] if self._rows else None def all(self) -> list[dict[str, Any]]: return list(self._rows) class FakeConnection: """A connection that records and then answers from the engine's queues.""" def __init__(self, engine: FakeStatusEngine) -> None: self._engine = engine def execute(self, statement: Any, parameters: dict[str, Any] | None = None) -> FakeResult: recorded = RecordedStatement(statement, parameters) self._engine.statements.append(recorded) if recorded.is_write: return self._engine._answer_write() return self._engine._answer_read() class _Scope: """The context manager `engine.begin()` and `engine.connect()` both hand out.""" def __init__(self, engine: FakeStatusEngine, transactional: bool) -> None: self._engine = engine self._transactional = transactional def __enter__(self) -> FakeConnection: self._engine.open_scopes += 1 if self._transactional: self._engine.transactions += 1 else: self._engine.connections += 1 return FakeConnection(self._engine) def __exit__(self, *_exc: object) -> bool: self._engine.open_scopes -= 1 return False class FakeStatusEngine: """An engine double that records statements and can be told to fail a write. Attributes: statements: every statement executed, in order. rowcounts: rowcounts for writes, consumed in order; the last one repeats. rows: rows for reads, consumed in order; the last one repeats. `None` means "no such row". write_errors: exceptions to raise instead of executing a write, consumed in order. A `None` entry means "this attempt succeeds", which is how a retry is modelled. transactions: how many `begin()` scopes were opened — one per write, or the writes are sharing a transaction. open_scopes: scopes currently open, so a test can assert none is left open. """ def __init__( self, rowcounts: list[int] | None = None, rows: list[dict[str, Any] | None] | None = None, write_errors: list[BaseException | None] | None = None, ) -> None: self.statements: list[RecordedStatement] = [] self.rowcounts: list[int] = list(rowcounts or []) self.rows: list[dict[str, Any] | None] = list(rows or []) self.write_errors: list[BaseException | None] = list(write_errors or []) self.transactions = 0 self.connections = 0 self.open_scopes = 0 self.disposed = False # ------------------------------------------------------------------ engine API def begin(self) -> _Scope: return _Scope(self, transactional=True) def connect(self) -> _Scope: return _Scope(self, transactional=False) def dispose(self) -> None: self.disposed = True # ------------------------------------------------------------------- answering def _answer_write(self) -> FakeResult: if self.write_errors: error = self.write_errors.pop(0) if error is not None: raise error rowcount = self.rowcounts.pop(0) if len(self.rowcounts) > 1 else (self.rowcounts or [1])[0] return FakeResult(rowcount=rowcount) def _answer_read(self) -> FakeResult: row = self.rows.pop(0) if len(self.rows) > 1 else (self.rows or [None])[0] return FakeResult(rows=[row] if row is not None else []) # ---------------------------------------------------------------- convenience @property def writes(self) -> list[RecordedStatement]: return [statement for statement in self.statements if statement.is_write] @property def reads(self) -> list[RecordedStatement]: return [statement for statement in self.statements if not statement.is_write] @property def tables(self) -> set[str]: return {statement.table for statement in self.statements if statement.table}