postgresql.py

"""
Plugin that imports the required fixtures to start the database
and adjusts the PostgreSQL "dbconnection" static config value.
"""
 
from contextlib import contextmanager
from enum import Enum
import typing
 
import pytest
 
from pytest_userver import sql
 
pytest_plugins = [
    'testsuite.databases.pgsql.pytest_plugin',
    'pytest_userver.plugins.core',
]
 
 
USERVER_CONFIG_HOOKS = ['userver_pg_config']
 
 
class RegisteredNtrxFailureType(Enum):
    Error = 1
    LogicError = 2
    RuntimeError = 3
    ConnectionError = 4
 
 
class RegisteredNtrx:
    def __init__(self, testpoint):
        self._registered_ntrx = dict()
        self._testpoint = testpoint
 
    def _enable_failure(
        self,
        name: str,
        failure_type: RegisteredNtrxFailureType,
    ) -> None:
        self._registered_ntrx[name] = failure_type
 
        @self._testpoint(f'pg_ntrx_execute::{name}')
        def _failure_tp(data):
            return {
                'inject_failure': self.is_failure_enabled(name),
                'failure_type': self._get_failure_type(name).name,
            }
 
    def _disable_failure(self, name: str) -> None:
        if self.is_failure_enabled(name):
            del self._registered_ntrx[name]
 
    def _get_failure_type(self, name: str) -> RegisteredNtrxFailureType:
        return self._registered_ntrx[name]
 
    def is_failure_enabled(self, name: str) -> bool:
        return name in self._registered_ntrx
 
    @contextmanager
    def mock_failure(
        self,
        name: str,
        failure_type: RegisteredNtrxFailureType = RegisteredNtrxFailureType.Error,
    ):
        self._enable_failure(name, failure_type)
        try:
            yield
        finally:
            self._disable_failure(name)
 
 
@pytest.fixture(scope='session')
def userver_pg_config(pgsql_local):
    """
    Returns a function that adjusts the static configuration file for
    the testsuite.
    Sets the `dbconnection` to the testsuite started PostgreSQL credentials
    if there's only one `dbconnection` in static config.
 
    @ingroup userver_testsuite_fixtures
    """
 
    if not pgsql_local:
        raise ValueError(
            'Override the "pgsql_local" fixture so that testsuite knowns how to start the PostgreSQL database',
        )
 
    if len(pgsql_local) > 1:
        raise ValueError(
            f'Found more than one entry in "pgsql_local": '
            f'{list(pgsql_local.keys())}. '
            f'The "userver_pg_config" fixture supports '
            f'only one entry in "pgsql_local" fixture. The '
            f'"userver_pg_config" fixture should be overridden and '
            f'the "dbconnection" for the components::Postgres '
            f'components should be adjusted via the overridden fixture.',
        )
 
    uri = list(pgsql_local.values())[0].get_uri()
 
    def _patch_config(config_yaml, config_vars):
        components = config_yaml['components_manager']['components']
        postgre_dbs = {
            name: params
            for name, params in components.items()
            if params and ('dbconnection' in params or 'dbconnection#env' in params)
        }
 
        if len(postgre_dbs) > 1:
            raise ValueError(
                f'Found more than one components with "dbconnection": '
                f'{list(postgre_dbs.keys())}. '
                f'The "userver_pg_config" fixture should be overridden and '
                f'the "dbconnection" for the components::Postgres '
                f'components should be adjusted via the overridden fixture.',
            )
 
        for config in postgre_dbs.values():
            config['dbconnection'] = uri
            config.pop('dbalias', None)
 
    return _patch_config
 
 
@pytest.fixture
def userver_pg_trx(
    testpoint,
) -> typing.Generator[sql.RegisteredTrx, None, None]:
    """
    The fixture maintains transaction fault injection state using
    RegisteredTrx class.
 
    @see pytest_userver.sql.RegisteredTrx
 
    @snippet postgresql/functional_tests/integration_tests/tests/test_trx_failure.py  fault injection
 
    @ingroup userver_testsuite_fixtures
    """  # noqa: E501
 
    registered = sql.RegisteredTrx()
 
    @testpoint('pg_trx_commit')
    def _pg_trx_tp(data):
        should_fail = registered.is_failure_enabled(data['trx_name'])
        return {'trx_should_fail': should_fail}
 
    yield registered
 
 
@pytest.fixture
def userver_pg_ntrx(testpoint) -> typing.Generator[RegisteredNtrx, None, None]:
    """
    The fixture maintains single query fault injection state using
    RegisteredNtrx class.
 
    @see pytest_userver.plugins.postgresql.RegisteredNtrx
 
    @snippet postgresql/functional_tests/integration_tests/tests/test_ntrx_failure.py  fault injection
 
    @ingroup userver_testsuite_fixtures
    """  # noqa: E501
 
    yield RegisteredNtrx(testpoint)