Skip to content

Commit b325829

Browse files
blink1073blink1073
committed
fix handling of client_knobs
1 parent 446cd96 commit b325829

File tree

10 files changed

+246
-139
lines changed

10 files changed

+246
-139
lines changed

‎test/__init__.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959

6060
sys.path[0:0] = [""]
6161

62+
from test.helpers import client_knobs, global_knobs
6263
from test.helpers_shared import (
6364
COMPRESSORS,
6465
IS_SRV,
@@ -67,10 +68,8 @@
6768
TEST_LOADBALANCER,
6869
TLS_OPTIONS,
6970
SystemCertsPatcher,
70-
client_knobs,
7171
db_pwd,
7272
db_user,
73-
global_knobs,
7473
host,
7574
is_server_resolvable,
7675
port,

‎test/asynchronous/__init__.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959

6060
sys.path[0:0] = [""]
6161

62+
from test.asynchronous.helpers import client_knobs, global_knobs
6263
from test.helpers_shared import (
6364
COMPRESSORS,
6465
IS_SRV,
@@ -67,10 +68,8 @@
6768
TEST_LOADBALANCER,
6869
TLS_OPTIONS,
6970
SystemCertsPatcher,
70-
client_knobs,
7171
db_pwd,
7272
db_user,
73-
global_knobs,
7473
host,
7574
is_server_resolvable,
7675
port,

‎test/asynchronous/helpers.py‎

Lines changed: 91 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,12 @@
1717

1818
import asyncio
1919
import threading
20-
from typing import Optional
20+
import traceback
21+
from functools import wraps
22+
from typing import Optional, no_type_check
2123

2224
from bson import SON
25+
from pymongo import common
2326
from pymongo._asyncio_task import create_task
2427
from pymongo.read_preferences import ReadPreference
2528

@@ -36,6 +39,93 @@ async def async_repl_set_step_down(client, **kwargs):
3639
await client.admin.command(cmd)
3740

3841

42+
class client_knobs:
43+
def __init__(
44+
self,
45+
heartbeat_frequency=None,
46+
min_heartbeat_interval=None,
47+
kill_cursor_frequency=None,
48+
events_queue_frequency=None,
49+
):
50+
self.heartbeat_frequency = heartbeat_frequency
51+
self.min_heartbeat_interval = min_heartbeat_interval
52+
self.kill_cursor_frequency = kill_cursor_frequency
53+
self.events_queue_frequency = events_queue_frequency
54+
55+
self.old_heartbeat_frequency = None
56+
self.old_min_heartbeat_interval = None
57+
self.old_kill_cursor_frequency = None
58+
self.old_events_queue_frequency = None
59+
self._enabled = False
60+
self._stack = None
61+
62+
def enable(self):
63+
self.old_heartbeat_frequency = common.HEARTBEAT_FREQUENCY
64+
self.old_min_heartbeat_interval = common.MIN_HEARTBEAT_INTERVAL
65+
self.old_kill_cursor_frequency = common.KILL_CURSOR_FREQUENCY
66+
self.old_events_queue_frequency = common.EVENTS_QUEUE_FREQUENCY
67+
68+
if self.heartbeat_frequency is not None:
69+
common.HEARTBEAT_FREQUENCY = self.heartbeat_frequency
70+
71+
if self.min_heartbeat_interval is not None:
72+
common.MIN_HEARTBEAT_INTERVAL = self.min_heartbeat_interval
73+
74+
if self.kill_cursor_frequency is not None:
75+
common.KILL_CURSOR_FREQUENCY = self.kill_cursor_frequency
76+
77+
if self.events_queue_frequency is not None:
78+
common.EVENTS_QUEUE_FREQUENCY = self.events_queue_frequency
79+
self._enabled = True
80+
# Store the allocation traceback to catch non-disabled client_knobs.
81+
self._stack = "".join(traceback.format_stack())
82+
83+
def __enter__(self):
84+
self.enable()
85+
86+
@no_type_check
87+
def disable(self):
88+
common.HEARTBEAT_FREQUENCY = self.old_heartbeat_frequency
89+
common.MIN_HEARTBEAT_INTERVAL = self.old_min_heartbeat_interval
90+
common.KILL_CURSOR_FREQUENCY = self.old_kill_cursor_frequency
91+
common.EVENTS_QUEUE_FREQUENCY = self.old_events_queue_frequency
92+
self._enabled = False
93+
94+
def __exit__(self, exc_type, exc_val, exc_tb):
95+
self.disable()
96+
97+
def __call__(self, func):
98+
def make_wrapper(f):
99+
@wraps(f)
100+
async def wrap(*args, **kwargs):
101+
with self:
102+
return await f(*args, **kwargs)
103+
104+
return wrap
105+
106+
return make_wrapper(func)
107+
108+
def __del__(self):
109+
if self._enabled:
110+
msg = (
111+
"ERROR: client_knobs still enabled! HEARTBEAT_FREQUENCY={}, "
112+
"MIN_HEARTBEAT_INTERVAL={}, KILL_CURSOR_FREQUENCY={}, "
113+
"EVENTS_QUEUE_FREQUENCY={}, stack:\n{}".format(
114+
common.HEARTBEAT_FREQUENCY,
115+
common.MIN_HEARTBEAT_INTERVAL,
116+
common.KILL_CURSOR_FREQUENCY,
117+
common.EVENTS_QUEUE_FREQUENCY,
118+
self._stack,
119+
)
120+
)
121+
self.disable()
122+
raise Exception(msg)
123+
124+
125+
# Global knobs to speed up the test suite.
126+
global_knobs = client_knobs(events_queue_frequency=0.05)
127+
128+
39129
if _IS_SYNC:
40130
PARENT = threading.Thread
41131
else:

‎test/asynchronous/test_retryable_writes.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
async_client_context,
3131
unittest,
3232
)
33-
from test.helpers_shared import client_knobs
33+
from test.asynchronous.helpers import client_knobs
3434
from test.utils_shared import (
3535
CMAPListener,
3636
DeprecationFilter,

‎test/asynchronous/test_session.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
async_client_context,
3636
unittest,
3737
)
38-
from test.helpers_shared import client_knobs
38+
from test.asynchronous.helpers import client_knobs
3939
from test.utils_shared import (
4040
EventListener,
4141
HeartbeatEventListener,

‎test/helpers.py‎

Lines changed: 91 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,12 @@
1717

1818
import asyncio
1919
import threading
20-
from typing import Optional
20+
import traceback
21+
from functools import wraps
22+
from typing import Optional, no_type_check
2123

2224
from bson import SON
25+
from pymongo import common
2326
from pymongo._asyncio_task import create_task
2427
from pymongo.read_preferences import ReadPreference
2528

@@ -36,6 +39,93 @@ def repl_set_step_down(client, **kwargs):
3639
client.admin.command(cmd)
3740

3841

42+
class client_knobs:
43+
def __init__(
44+
self,
45+
heartbeat_frequency=None,
46+
min_heartbeat_interval=None,
47+
kill_cursor_frequency=None,
48+
events_queue_frequency=None,
49+
):
50+
self.heartbeat_frequency = heartbeat_frequency
51+
self.min_heartbeat_interval = min_heartbeat_interval
52+
self.kill_cursor_frequency = kill_cursor_frequency
53+
self.events_queue_frequency = events_queue_frequency
54+
55+
self.old_heartbeat_frequency = None
56+
self.old_min_heartbeat_interval = None
57+
self.old_kill_cursor_frequency = None
58+
self.old_events_queue_frequency = None
59+
self._enabled = False
60+
self._stack = None
61+
62+
def enable(self):
63+
self.old_heartbeat_frequency = common.HEARTBEAT_FREQUENCY
64+
self.old_min_heartbeat_interval = common.MIN_HEARTBEAT_INTERVAL
65+
self.old_kill_cursor_frequency = common.KILL_CURSOR_FREQUENCY
66+
self.old_events_queue_frequency = common.EVENTS_QUEUE_FREQUENCY
67+
68+
if self.heartbeat_frequency is not None:
69+
common.HEARTBEAT_FREQUENCY = self.heartbeat_frequency
70+
71+
if self.min_heartbeat_interval is not None:
72+
common.MIN_HEARTBEAT_INTERVAL = self.min_heartbeat_interval
73+
74+
if self.kill_cursor_frequency is not None:
75+
common.KILL_CURSOR_FREQUENCY = self.kill_cursor_frequency
76+
77+
if self.events_queue_frequency is not None:
78+
common.EVENTS_QUEUE_FREQUENCY = self.events_queue_frequency
79+
self._enabled = True
80+
# Store the allocation traceback to catch non-disabled client_knobs.
81+
self._stack = "".join(traceback.format_stack())
82+
83+
def __enter__(self):
84+
self.enable()
85+
86+
@no_type_check
87+
def disable(self):
88+
common.HEARTBEAT_FREQUENCY = self.old_heartbeat_frequency
89+
common.MIN_HEARTBEAT_INTERVAL = self.old_min_heartbeat_interval
90+
common.KILL_CURSOR_FREQUENCY = self.old_kill_cursor_frequency
91+
common.EVENTS_QUEUE_FREQUENCY = self.old_events_queue_frequency
92+
self._enabled = False
93+
94+
def __exit__(self, exc_type, exc_val, exc_tb):
95+
self.disable()
96+
97+
def __call__(self, func):
98+
def make_wrapper(f):
99+
@wraps(f)
100+
def wrap(*args, **kwargs):
101+
with self:
102+
return f(*args, **kwargs)
103+
104+
return wrap
105+
106+
return make_wrapper(func)
107+
108+
def __del__(self):
109+
if self._enabled:
110+
msg = (
111+
"ERROR: client_knobs still enabled! HEARTBEAT_FREQUENCY={}, "
112+
"MIN_HEARTBEAT_INTERVAL={}, KILL_CURSOR_FREQUENCY={}, "
113+
"EVENTS_QUEUE_FREQUENCY={}, stack:\n{}".format(
114+
common.HEARTBEAT_FREQUENCY,
115+
common.MIN_HEARTBEAT_INTERVAL,
116+
common.KILL_CURSOR_FREQUENCY,
117+
common.EVENTS_QUEUE_FREQUENCY,
118+
self._stack,
119+
)
120+
)
121+
self.disable()
122+
raise Exception(msg)
123+
124+
125+
# Global knobs to speed up the test suite.
126+
global_knobs = client_knobs(events_queue_frequency=0.05)
127+
128+
39129
if _IS_SYNC:
40130
PARENT = threading.Thread
41131
else:

‎test/helpers_shared.py‎

Lines changed: 1 addition & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
from unittest import SkipTest
3636

3737
from bson.son import SON
38-
from pymongo import common, message
38+
from pymongo import message
3939
from pymongo.ssl_support import HAVE_SSL, _ssl # type:ignore[attr-defined]
4040
from pymongo.synchronous.uri_parser import parse_uri
4141

@@ -148,89 +148,6 @@ def _create_user(authdb, user, pwd=None, roles=None, **kwargs):
148148
return authdb.command(cmd)
149149

150150

151-
class client_knobs:
152-
def __init__(
153-
self,
154-
heartbeat_frequency=None,
155-
min_heartbeat_interval=None,
156-
kill_cursor_frequency=None,
157-
events_queue_frequency=None,
158-
):
159-
self.heartbeat_frequency = heartbeat_frequency
160-
self.min_heartbeat_interval = min_heartbeat_interval
161-
self.kill_cursor_frequency = kill_cursor_frequency
162-
self.events_queue_frequency = events_queue_frequency
163-
164-
self.old_heartbeat_frequency = None
165-
self.old_min_heartbeat_interval = None
166-
self.old_kill_cursor_frequency = None
167-
self.old_events_queue_frequency = None
168-
self._enabled = False
169-
self._stack = None
170-
171-
def enable(self):
172-
self.old_heartbeat_frequency = common.HEARTBEAT_FREQUENCY
173-
self.old_min_heartbeat_interval = common.MIN_HEARTBEAT_INTERVAL
174-
self.old_kill_cursor_frequency = common.KILL_CURSOR_FREQUENCY
175-
self.old_events_queue_frequency = common.EVENTS_QUEUE_FREQUENCY
176-
177-
if self.heartbeat_frequency is not None:
178-
common.HEARTBEAT_FREQUENCY = self.heartbeat_frequency
179-
180-
if self.min_heartbeat_interval is not None:
181-
common.MIN_HEARTBEAT_INTERVAL = self.min_heartbeat_interval
182-
183-
if self.kill_cursor_frequency is not None:
184-
common.KILL_CURSOR_FREQUENCY = self.kill_cursor_frequency
185-
186-
if self.events_queue_frequency is not None:
187-
common.EVENTS_QUEUE_FREQUENCY = self.events_queue_frequency
188-
self._enabled = True
189-
# Store the allocation traceback to catch non-disabled client_knobs.
190-
self._stack = "".join(traceback.format_stack())
191-
192-
def __enter__(self):
193-
self.enable()
194-
195-
@no_type_check
196-
def disable(self):
197-
common.HEARTBEAT_FREQUENCY = self.old_heartbeat_frequency
198-
common.MIN_HEARTBEAT_INTERVAL = self.old_min_heartbeat_interval
199-
common.KILL_CURSOR_FREQUENCY = self.old_kill_cursor_frequency
200-
common.EVENTS_QUEUE_FREQUENCY = self.old_events_queue_frequency
201-
self._enabled = False
202-
203-
def __exit__(self, exc_type, exc_val, exc_tb):
204-
self.disable()
205-
206-
def __call__(self, func):
207-
def make_wrapper(f):
208-
@wraps(f)
209-
async def wrap(*args, **kwargs):
210-
with self:
211-
return await f(*args, **kwargs)
212-
213-
return wrap
214-
215-
return make_wrapper(func)
216-
217-
def __del__(self):
218-
if self._enabled:
219-
msg = (
220-
"ERROR: client_knobs still enabled! HEARTBEAT_FREQUENCY={}, "
221-
"MIN_HEARTBEAT_INTERVAL={}, KILL_CURSOR_FREQUENCY={}, "
222-
"EVENTS_QUEUE_FREQUENCY={}, stack:\n{}".format(
223-
common.HEARTBEAT_FREQUENCY,
224-
common.MIN_HEARTBEAT_INTERVAL,
225-
common.KILL_CURSOR_FREQUENCY,
226-
common.EVENTS_QUEUE_FREQUENCY,
227-
self._stack,
228-
)
229-
)
230-
self.disable()
231-
raise Exception(msg)
232-
233-
234151
def _all_users(db):
235152
return {u["user"] for u in db.command("usersInfo").get("users", [])}
236153

@@ -287,10 +204,6 @@ def print_thread_stacks(pid: int) -> None:
287204
sys.stderr.write(res.stdout)
288205

289206

290-
# Global knobs to speed up the test suite.
291-
global_knobs = client_knobs(events_queue_frequency=0.05)
292-
293-
294207
def _get_executors(topology):
295208
executors = []
296209
for server in topology._servers.values():

‎test/test_retryable_writes.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
client_context,
3131
unittest,
3232
)
33-
from test.helpers_shared import client_knobs
33+
from test.helpers import client_knobs
3434
from test.utils_shared import (
3535
CMAPListener,
3636
DeprecationFilter,

‎test/test_session.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
client_context,
3636
unittest,
3737
)
38-
from test.helpers_shared import client_knobs
38+
from test.helpers import client_knobs
3939
from test.utils_shared import (
4040
EventListener,
4141
HeartbeatEventListener,

0 commit comments

Comments
 (0)