userver: /data/code/userver/testsuite/pytest_plugins/pytest_userver/plugins/postgresql.py Source File
Loading...
Searching...
No Matches
postgresql.py
1"""
2Plugin that imports the required fixtures to start the database
3and adjusts the PostgreSQL "dbconnection" static config value.
4"""
5
6from contextlib import contextmanager
7import typing
8
9import pytest
10
11from pytest_userver import sql
12
13
14pytest_plugins = [
15 'testsuite.databases.pgsql.pytest_plugin',
16 'pytest_userver.plugins.core',
17]
18
19
20USERVER_CONFIG_HOOKS = ['userver_pg_config']
21
22
24 def __init__(self, testpoint):
25 self._registered_ntrx = set()
26 self._testpoint = testpoint
27
28 def _enable_failure(self, name: str) -> None:
29 self._registered_ntrx.add(name)
30
31 @self._testpoint(f'pg_ntrx_execute::{name}')
32 def _failure_tp(data):
33 return {'inject_failure': self.is_failure_enabled(name)}
34
35 def _disable_failure(self, name: str) -> None:
36 if self.is_failure_enabled(name):
37 self._registered_ntrx.remove(name)
38
39 def is_failure_enabled(self, name: str) -> bool:
40 return name in self._registered_ntrx
41
42 @contextmanager
43 def mock_failure(self, name: str):
44 self._enable_failure(name)
45 try:
46 yield
47 finally:
48 self._disable_failure(name)
49
50
51@pytest.fixture(scope='session')
52def userver_pg_config(pgsql_local):
53 """
54 Returns a function that adjusts the static configuration file for
55 the testsuite.
56 Sets the `dbconnection` to the testsuite started PostgreSQL credentials
57 if there's only one `dbconnection` in static config.
58
59 @ingroup userver_testsuite_fixtures
60 """
61
62 if not pgsql_local:
63 raise ValueError(
64 'Override the "pgsql_local" fixture so that testsuite knowns how '
65 'to start the PostgreSQL database',
66 )
67
68 if len(pgsql_local) > 1:
69 raise ValueError(
70 f'Found more than one entry in "pgsql_local": '
71 f'{list(pgsql_local.keys())}. '
72 f'The "userver_pg_config" fixture supports '
73 f'only one entry in "pgsql_local" fixture. The '
74 f'"userver_pg_config" fixture should be overridden and '
75 f'the "dbconnection" for the components::Postgres '
76 f'components should be adjusted via the overridden fixture.',
77 )
78
79 uri = list(pgsql_local.values())[0].get_uri()
80
81 def _patch_config(config_yaml, config_vars):
82 components = config_yaml['components_manager']['components']
83 postgre_dbs = {
84 name: params
85 for name, params in components.items()
86 if params
87 and ('dbconnection' in params or 'dbconnection#env' in params)
88 }
89
90 if len(postgre_dbs) > 1:
91 raise ValueError(
92 f'Found more than one components with "dbconnection": '
93 f'{list(postgre_dbs.keys())}. '
94 f'The "userver_pg_config" fixture should be overridden and '
95 f'the "dbconnection" for the components::Postgres '
96 f'components should be adjusted via the overridden fixture.',
97 )
98
99 for config in postgre_dbs.values():
100 config['dbconnection'] = uri
101 config.pop('dbalias', None)
102
103 return _patch_config
104
105
106@pytest.fixture
108 testpoint,
109) -> typing.Generator[sql.RegisteredTrx, None, None]:
110 """
111 The fixture maintains transaction fault injection state using
112 RegisteredTrx class.
113
114 @see pytest_userver.sql.RegisteredTrx
115
116 @snippet postgresql/functional_tests/integration_tests/tests/test_trx_failure.py fault injection
117
118 @ingroup userver_testsuite_fixtures
119 """ # noqa: E501
120
121 registered = sql.RegisteredTrx()
122
123 @testpoint('pg_trx_commit')
124 def _pg_trx_tp(data):
125 should_fail = registered.is_failure_enabled(data['trx_name'])
126 return {'trx_should_fail': should_fail}
127
128 yield registered
129
130
131@pytest.fixture
132def userver_pg_ntrx(testpoint) -> typing.Generator[RegisteredNtrx, None, None]:
133 """
134 The fixture maintains single query fault injection state using
135 RegisteredNtrx class.
136
137 @see pytest_userver.plugins.postgresql.RegisteredNtrx
138
139 @snippet postgresql/functional_tests/integration_tests/tests/test_ntrx_failure.py fault injection
140
141 @ingroup userver_testsuite_fixtures
142 """ # noqa: E501
143
144 yield RegisteredNtrx(testpoint)