1"""
  2Database test lifecycle, registered under the `plain.testing` entry point.
  3
  4Every test runs against a dedicated test database (created once per run,
  5as a clone of a migrated and converged template that stays between runs)
  6inside a transaction that rolls back afterward.
  7Tests tagged with `@isolated_db` get their own separately-created database
  8for the duration of the test instead of a rolled-back transaction — for
  9DDL-heavy tests (migrations, convergence) that can't run inside a
 10transaction that never commits.
 11
 12The databases are the run's own, and while one is in use it is the
 13configured database for the whole process: see `database.py`.
 14
 15The rolled-back tests share one connection, opened by the first of them.
 16The rollback undoes everything a test did in its transaction. What belongs
 17to the session and not to a transaction carries over to the next test: a
 18session-level advisory lock, a `LISTEN`, a server-side prepared statement.
 19A test that needs a session of its own is `@isolated_db`.
 20"""
 21
 22from collections.abc import Generator
 23from contextlib import ExitStack, contextmanager
 24
 25from plain.runtime import settings
 26from plain.testing import CollectedTest, TestLifecycle
 27
 28from .. import transaction
 29from ..base import Model
 30from ..db import get_connection
 31from ..otel import suppress_db_tracing
 32from ..query import QuerySet
 33from .database import RunDatabases, use_test_database
 34from .decorators import ISOLATED_DB_TAG
 35from .schema import describe_schema
 36
 37# A model instance that prints in less than this is printed on one line.
 38_ONE_LINE = 80
 39
 40
 41def describe_model_instance(instance: Model) -> str:
 42    """
 43    A model instance with its fields: `Widget(id=1, name='bolt', size='m')`,
 44    or one field to a line when that is long. Its `repr` is its class and
 45    its id.
 46
 47    A field that was deferred is read from where a loaded one is kept, so
 48    it is reported as not loaded and is not fetched.
 49
 50    A field whose type says its value is a secret (`value_is_secret`: an
 51    encrypted field, a password) is printed as withheld. One with nothing
 52    in it says so, since that gives nothing away.
 53    """
 54    fields = []
 55    for field in instance._model_meta.fields:
 56        if field.name not in instance.__dict__:
 57            fields.append(f"{field.name}=<not loaded>")
 58            continue
 59        value = instance.__dict__[field.name]
 60        if field.value_is_secret and value not in (None, ""):
 61            fields.append(f"{field.name}=<withheld>")
 62        else:
 63            fields.append(f"{field.name}={value!r}")
 64
 65    name = type(instance).__name__
 66    on_one_line = f"{name}({', '.join(fields)})"
 67    if len(on_one_line) <= _ONE_LINE:
 68        return on_one_line
 69    return "\n".join([f"{name}(", *(f"    {field}," for field in fields), ")"])
 70
 71
 72def describe_queryset(queryset: QuerySet) -> str:
 73    """
 74    A queryset by what it holds. One that has run prints its rows. One that
 75    hasn't prints the SQL it would run, and doesn't run it.
 76    """
 77    name = f"{type(queryset).__name__} of {queryset.model.__name__}"
 78    rows = queryset._result_cache
 79    if rows is not None:
 80        count = "1 row" if len(rows) == 1 else f"{len(rows):,} rows"
 81        return f"<{name}, {count}: {rows!r}>"
 82    try:
 83        # Compiled from a copy. Compiling settles a query's joins and
 84        # aliases, which is a change to the test's own queryset.
 85        sql, params = queryset._chain().sql_query.sql_with_params()
 86    except Exception:
 87        return f"<{name}, not run>"
 88    return f"<{name}, not run: {sql} with {params!r}>"
 89
 90
 91class PostgresTestLifecycle(TestLifecycle):
 92    required_package = "plain.postgres"
 93
 94    def __init__(self) -> None:
 95        # Holds the run's claim and its shared database, from setup to teardown.
 96        self._test_database = ExitStack()
 97        self._run: RunDatabases | None = None
 98
 99    def setup_worker(self) -> None:
100        run = RunDatabases(
101            runtime_url=str(settings.POSTGRES_URL),
102            management_url=str(settings.POSTGRES_MANAGEMENT_URL),
103        )
104        self._run = run
105        with suppress_db_tracing():
106            run.claim(schema=describe_schema(server_version=run.server_version()))
107            self._test_database.callback(run.release)
108            self._test_database.enter_context(
109                use_test_database(run=run, name=run.shared_name, verbosity=0)
110            )
111
112    def teardown_worker(self) -> None:
113        with suppress_db_tracing():
114            # Closes the connection the tests shared, drops the database,
115            # then gives up the run's claim.
116            self._test_database.close()
117
118    def describe_setup(self) -> tuple[tuple[str, float], ...]:
119        # Whether the template was built by this run or cloned, and how long
120        # each took.
121        assert self._run is not None
122        return tuple(self._run.setup_parts)
123
124    @contextmanager
125    def around_test(self, test: CollectedTest) -> Generator[None]:
126        if ISOLATED_DB_TAG in test.tags:
127            yield from self._run_in_isolated_database(test)
128        else:
129            yield from self._run_in_rolled_back_transaction()
130
131    def describe_value(self, value: object) -> str | None:
132        if isinstance(value, Model):
133            return describe_model_instance(value)
134        if isinstance(value, QuerySet):
135            return describe_queryset(value)
136        return None
137
138    def _run_in_rolled_back_transaction(self) -> Generator[None]:
139        with suppress_db_tracing():
140            atomic = transaction.atomic()
141            atomic._from_testcase = True
142            atomic.__enter__()
143            # The transaction isn't open yet: the driver begins it with the
144            # first statement. Until then, code that opens psycopg's own
145            # `connection.transaction()` finds none, begins one itself, and
146            # commits it, and what it did outlives the test. So send the
147            # first statement here. The connection was the last test's too,
148            # and a session-level advisory lock is the session's, which a
149            # rollback doesn't release: the statement gives back any that
150            # test left held.
151            with get_connection().cursor() as cursor:
152                cursor.execute("SELECT pg_advisory_unlock_all()")
153
154        try:
155            yield
156        finally:
157            with suppress_db_tracing():
158                # Constraints are never deferred (foreign keys are NOT
159                # DEFERRABLE), so there is nothing to check before rolling
160                # back — the database has already rejected any violation.
161                get_connection().set_rollback(True)
162                atomic.__exit__(None, None, None)
163                # The connection stays open for the next test. Opening one
164                # costs a hundred times what the rollback does.
165
166    def _run_in_isolated_database(self, test: CollectedTest) -> Generator[None]:
167        run = self._run
168        assert run is not None
169        test_name = test.id.rpartition("::")[2]
170
171        ctx = use_test_database(run=run, name=run.isolated_name(test_name), verbosity=0)
172        with suppress_db_tracing():
173            ctx.__enter__()
174        try:
175            yield
176        finally:
177            with suppress_db_tracing():
178                ctx.__exit__(None, None, None)