sientia-dataops-laborious-legacy/tests/helpers/status_engine.py

"""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}