userver: /home/antonyzhilin/arcadia/taxi/uservices/userver/testsuite/pytest_plugins/pytest_userver/plugins/postgresql.py Source File
All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Modules Pages Concepts
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
7from enum import Enum
8import typing
9
10import pytest
11
12from pytest_userver import sql
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 Error = 1
25 LogicError = 2
26 RuntimeError = 3
27 ConnectionError = 4
28
29
31 def __init__(self, testpoint):
32 self._registered_ntrx = dict()
33 self._testpoint = testpoint
34
35 def _enable_failure(
36 self,
37 name: str,
38 failure_type: RegisteredNtrxFailureType,
39 ) -> None:
40 self._registered_ntrx[name] = failure_type
41
42 @self._testpoint(f'pg_ntrx_execute::{name}')
43 def _failure_tp(data):
44 return {
45 'inject_failure': self.is_failure_enabled(name),
46 'failure_type': self._get_failure_type(name).name,
47 }
48
49 def _disable_failure(self, name: str) -> None:
50 if self.is_failure_enabled(name):
51 del self._registered_ntrx[name]
52
53 def _get_failure_type(self, name: str) -> RegisteredNtrxFailureType:
54 return self._registered_ntrx[name]
55
56 def is_failure_enabled(self, name: str) -> bool:
57 return name in self._registered_ntrx
58
59 @contextmanager
60 def mock_failure(
61 self,
62 name: str,
63 failure_type: RegisteredNtrxFailureType = RegisteredNtrxFailureType.Error,
64 ):
65 self._enable_failure(name, failure_type)
66 try:
67 yield
68 finally:
69 self._disable_failure(name)
70
71
72@pytest.fixture(scope='session')
73def userver_pg_config(pgsql_local):
74 """
75 Returns a function that adjusts the static configuration file for
76 the testsuite.
77 Sets the `dbconnection` to the testsuite started PostgreSQL credentials
78 if there's only one `dbconnection` in static config.
79
80 @ingroup userver_testsuite_fixtures
81 """
82
83 if not pgsql_local:
84 raise ValueError(
85 'Override the "pgsql_local" fixture so that testsuite knowns how to start the PostgreSQL database',
86 )
87
88 if len(pgsql_local) > 1:
89 raise ValueError(
90 f'Found more than one entry in "pgsql_local": '
91 f'{list(pgsql_local.keys())}. '
92 f'The "userver_pg_config" fixture supports '
93 f'only one entry in "pgsql_local" fixture. The '
94 f'"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 uri = list(pgsql_local.values())[0].get_uri()
100
101 def _patch_config(config_yaml, config_vars):
102 components = config_yaml['components_manager']['components']
103 postgre_dbs = {
104 name: params
105 for name, params in components.items()
106 if params and ('dbconnection' in params or 'dbconnection#env' in params)
107 }
108
109 if len(postgre_dbs) > 1:
110 raise ValueError(
111 f'Found more than one components with "dbconnection": '
112 f'{list(postgre_dbs.keys())}. '
113 f'The "userver_pg_config" fixture should be overridden and '
114 f'the "dbconnection" for the components::Postgres '
115 f'components should be adjusted via the overridden fixture.',
116 )
117
118 for config in postgre_dbs.values():
119 config['dbconnection'] = uri
120 config.pop('dbalias', None)
121
122 return _patch_config
123
124
125@pytest.fixture
127 testpoint,
128) -> typing.Generator[sql.RegisteredTrx, None, None]:
129 """
130 The fixture maintains transaction fault injection state using
131 RegisteredTrx class.
132
133 @see pytest_userver.sql.RegisteredTrx
134
135 @snippet postgresql/functional_tests/integration_tests/tests/test_trx_failure.py fault injection
136
137 @ingroup userver_testsuite_fixtures
138 """ # noqa: E501
139
140 registered = sql.RegisteredTrx()
141
142 @testpoint('pg_trx_commit')
143 def _pg_trx_tp(data):
144 should_fail = registered.is_failure_enabled(data['trx_name'])
145 return {'trx_should_fail': should_fail}
146
147 yield registered
148
149
150@pytest.fixture
151def userver_pg_ntrx(testpoint) -> typing.Generator[RegisteredNtrx, None, None]:
152 """
153 The fixture maintains single query fault injection state using
154 RegisteredNtrx class.
155
156 @see pytest_userver.plugins.postgresql.RegisteredNtrx
157
158 @snippet postgresql/functional_tests/integration_tests/tests/test_ntrx_failure.py fault injection
159
160 @ingroup userver_testsuite_fixtures
161 """ # noqa: E501
162
163 yield RegisteredNtrx(testpoint)