"""Test helpers a brick imports instead of hand-rolling them.
`count_statements(engine)` counts what a block sends to the database: every statement, and the ROUND TRIPS
it cost — the number that matters on the pg wire (one per statement, except that a psycopg 3 pipeline sends
all of its statements in one; see erp_platform.repository). A brick's statement-budget test:
from erp_platform.testing import count_statements
def test_confirm_stays_within_its_budget(engine):
with count_statements(engine) as n:
must(api("POST", f"/orders/{o}/confirm", REP, {}))
assert n.statements <= 12, n.statements
`n` also compares as its statement count (`n <= 12`, `int(n)`), and keeps the SQL in `n.sql` for the message
of a failing budget."""
from contextlib import contextmanager
from dataclasses import dataclass, field
from sqlalchemy import event
from erp_platform import repository
@dataclass
class Counted:
statements: int = 0
round_trips: int = 0
sql: list = field(default_factory=list)
def __int__(self) -> int:
return self.statements
def __index__(self) -> int:
return self.statements
def __eq__(self, other):
return self.statements == other if isinstance(other, int) else NotImplemented
def __lt__(self, other):
return self.statements < int(other)
def __le__(self, other):
return self.statements <= int(other)
def __gt__(self, other):
return self.statements > int(other)
def __ge__(self, other):
return self.statements >= int(other)
__hash__ = None
@contextmanager
def count_statements(engine):
"""Count the statements (and round trips) sent on `engine` inside the block, from any connection or
thread. A pipelined load is counted as its statements and ONE round trip."""
box = Counted()
def _mine(conn) -> bool:
return getattr(conn, "engine", None) is engine
def _cursor(conn, cursor, statement, parameters, context, executemany):
box.statements += 1
box.sql.append(statement)
if not conn.info.get("erp_pipeline"):
box.round_trips += 1
def _pipeline(conn, statements):
if _mine(conn):
box.statements += statements
if not (statements and conn.info.get("erp_pipeline")): # a raw statement queued INSIDE a pipeline: no trip of its own
box.round_trips += 1
if statements:
box.sql.append(f"-- pipeline of {statements} reads")
event.listen(engine, "before_cursor_execute", _cursor)
repository.PIPELINE_LISTENERS.append(_pipeline)
try:
yield box
finally:
event.remove(engine, "before_cursor_execute", _cursor)
repository.PIPELINE_LISTENERS.remove(_pipeline)