Skip to content

Commit d41111a

Browse files
authored
test: isolate retry HTTP mocks from background telemetry (#1055)
* test: isolate retry HTTP mocks from background telemetry Signed-off-by: Vu Anh Phung <[email protected]> * test: keep mocked connections out of shared HTTP pools Signed-off-by: Vu Anh Phung <[email protected]> * test: isolate pool cleanup on transport failures Signed-off-by: Vu Anh Phung <[email protected]> --------- Signed-off-by: Vu Anh Phung <[email protected]>
1 parent 0c1b9e8 commit d41111a

2 files changed

Lines changed: 179 additions & 126 deletions

File tree

‎tests/e2e/common/retry_test_mixins.py‎

Lines changed: 40 additions & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
1-
from contextlib import contextmanager
1+
from contextlib import contextmanager, ExitStack
2+
import threading
23
import time
34
from typing import Optional, List
45
from unittest.mock import MagicMock, PropertyMock, patch
56
import io
67

78
import pytest
9+
from urllib3.connectionpool import HTTPSConnectionPool
810
from urllib3.exceptions import MaxRetryError
911

1012
from databricks.sql.auth.retry import DatabricksRetryPolicy
@@ -15,97 +17,35 @@
1517
SessionAlreadyClosedError,
1618
UnsafeToRetryError,
1719
)
18-
from databricks.sql.telemetry.telemetry_client import (
19-
NoopTelemetryClient,
20-
TelemetryClient,
21-
TelemetryClientFactory,
22-
_TelemetryClientHolder,
23-
)
2420

2521

2622
@contextmanager
27-
def _isolated_from_telemetry():
28-
# Tests that mock urllib3 globally (via _get_conn / _validate_conn) also
29-
# intercept background telemetry pushes from the shared
30-
# TelemetryClientFactory executor — inflating mock.call_count and breaking
31-
# assertions like `call_count == 6`. Three layers of defence:
32-
#
33-
# 1. Drain TelemetryClientFactory and override initialize_telemetry_client
34-
# so new connections install NoopTelemetryClient (which submits nothing).
35-
# 2. Patch TelemetryClient._send_telemetry to a no-op as a backstop — covers
36-
# any real TelemetryClient instance that slips in (e.g. a stale module-
37-
# global, a code path that bypasses initialize_telemetry_client, or
38-
# anything created before this context entered).
39-
# 3. Patch TelemetryClient._export_event to a no-op so even if a real
40-
# client receives an event, the event never reaches the queue and the
41-
# flush logic never fires.
42-
#
43-
# Without layer 2/3 we have observed `/telemetry-ext` requests showing up
44-
# in merge_group runs (concurrent CI load on the warehouse stresses paths
45-
# that single-test runs don't hit), inflating retry-test counts and
46-
# producing flaky AssertionErrors.
47-
with TelemetryClientFactory._lock:
48-
saved_clients = TelemetryClientFactory._clients
49-
saved_executor = TelemetryClientFactory._executor
50-
saved_initialized = TelemetryClientFactory._initialized
51-
for holder in list(saved_clients.values()):
52-
try:
53-
holder.client.close()
54-
except Exception:
55-
pass
56-
if saved_executor is not None:
57-
try:
58-
TelemetryClientFactory._stop_flush_thread()
59-
saved_executor.shutdown(wait=True)
60-
except Exception:
61-
pass
62-
TelemetryClientFactory._clients = {}
63-
TelemetryClientFactory._executor = None
64-
TelemetryClientFactory._initialized = False
65-
66-
def _install_noop(*args, host_url=None, **kwargs):
67-
host_url = TelemetryClientFactory.getHostUrlSafely(host_url)
68-
with TelemetryClientFactory._lock:
69-
TelemetryClientFactory._clients[host_url] = _TelemetryClientHolder(
70-
NoopTelemetryClient()
23+
def _mock_pool_method(method_name):
24+
"""Mock the calling test thread; worker threads use the original pool methods."""
25+
original = getattr(HTTPSConnectionPool, method_name)
26+
test_thread = threading.current_thread()
27+
mock_method = MagicMock()
28+
29+
def dispatch(pool, *args, **kwargs):
30+
if threading.current_thread() is test_thread:
31+
return mock_method(*args, **kwargs)
32+
return original(pool, *args, **kwargs)
33+
34+
with ExitStack() as stack:
35+
stack.enter_context(patch.object(HTTPSConnectionPool, method_name, dispatch))
36+
if method_name == "_get_conn":
37+
mock_method.return_value.close.return_value = None
38+
original_put = HTTPSConnectionPool._put_conn
39+
40+
def put_connection(pool, connection):
41+
# Fake acquisition never removed a slot from the real queue.
42+
if threading.current_thread() is not test_thread:
43+
return original_put(pool, connection)
44+
45+
stack.enter_context(
46+
patch.object(HTTPSConnectionPool, "_put_conn", put_connection)
7147
)
72-
73-
def _noop_send(self, events):
74-
return None
75-
76-
def _noop_export(self, event):
77-
return None
78-
79-
try:
80-
def _noop_send_with_unified_client(self, *args, **kwargs):
81-
return None
82-
83-
with patch.object(
84-
TelemetryClientFactory,
85-
"initialize_telemetry_client",
86-
staticmethod(_install_noop),
87-
), patch.object(
88-
TelemetryClient, "_send_telemetry", _noop_send
89-
), patch.object(
90-
TelemetryClient, "_export_event", _noop_export
91-
), patch.object(
92-
# Backstop layer 4: no-op the actual HTTP boundary. _send_telemetry
93-
# only suppresses NEW submissions; a future already submitted (e.g. by
94-
# the holder.client.close() flush above, which runs before these
95-
# patches apply) still calls _send_with_unified_client on a background
96-
# thread and would hit the network inside the test's urllib3 mock.
97-
# Patching the HTTP method itself closes that already-submitted-future
98-
# gap for any caller of this helper.
99-
TelemetryClient,
100-
"_send_with_unified_client",
101-
_noop_send_with_unified_client,
102-
):
103-
yield
104-
finally:
105-
with TelemetryClientFactory._lock:
106-
TelemetryClientFactory._clients = saved_clients
107-
TelemetryClientFactory._executor = saved_executor
108-
TelemetryClientFactory._initialized = saved_initialized
48+
yield mock_method
10949

11050

11151
class Client429ResponseMixin:
@@ -260,7 +200,7 @@ def mocked_server_response(
260200

261201
mock_response = SimpleHttpResponse(status, headers, redirect_location)
262202

263-
with patch("urllib3.connectionpool.HTTPSConnectionPool._get_conn") as getconn_mock:
203+
with _mock_pool_method("_get_conn") as getconn_mock:
264204
getconn_mock.return_value.getresponse.return_value = mock_response
265205
try:
266206
yield getconn_mock
@@ -288,7 +228,7 @@ def mock_sequential_server_responses(responses: List[dict]):
288228
for resp in responses
289229
]
290230

291-
with patch("urllib3.connectionpool.HTTPSConnectionPool._get_conn") as getconn_mock:
231+
with _mock_pool_method("_get_conn") as getconn_mock:
292232
getconn_mock.return_value.getresponse.side_effect = mock_responses
293233
try:
294234
yield getconn_mock
@@ -306,21 +246,7 @@ class PySQLRetryTestsMixin:
306246
"_retry_stop_after_attempts_count": 5,
307247
"_retry_stop_after_attempts_duration": 30,
308248
"_retry_delay_default": 0.5,
309-
# Disable telemetry on every retry-count test. These tests patch urllib3
310-
# globally (HTTPSConnectionPool._get_conn / _validate_conn) and assert the
311-
# connection's own request was retried exactly N times. With telemetry on
312-
# (the default), the same connection makes two kinds of side-channel HTTP
313-
# calls through the same mocked pool that inflate the count and flake the
314-
# assertion:
315-
# 1. A synchronous feature-flag GET during Connection.__init__
316-
# (is_telemetry_enabled -> FeatureFlagsContext._refresh_flags), and
317-
# 2. Async telemetry POSTs to /telemetry-ext from the background executor
318-
# and the periodic flush thread.
319-
# enable_telemetry=False short-circuits is_telemetry_enabled before the
320-
# feature-flag fetch and installs a NoopTelemetryClient, removing both
321-
# sources at their origin without changing the retry behavior under test.
322-
# Every test in this mixin merges _retry_policy into its connection params,
323-
# so setting it here covers them all. See also _isolated_from_telemetry().
249+
# Avoid synchronous feature-flag requests during connection setup.
324250
"enable_telemetry": False,
325251
}
326252

@@ -331,10 +257,7 @@ class PySQLRetryTestsMixin:
331257
{"use_sea": True},
332258
],
333259
)
334-
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
335-
def test_retry_urllib3_settings_are_honored(
336-
self, mock_send_telemetry, extra_params
337-
):
260+
def test_retry_urllib3_settings_are_honored(self, extra_params):
338261
"""Databricks overrides some of urllib3's configuration. This tests confirms that what configuration
339262
we DON'T override is preserved in urllib3's internals
340263
"""
@@ -361,12 +284,9 @@ def test_retry_urllib3_settings_are_honored(
361284
{"use_sea": True},
362285
],
363286
)
364-
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
365-
def test_oserror_retries(self, mock_send_telemetry, extra_params):
287+
def test_oserror_retries(self, extra_params):
366288
"""If a network error occurs during make_request, the request is retried according to policy"""
367-
with _isolated_from_telemetry(), patch(
368-
"urllib3.connectionpool.HTTPSConnectionPool._validate_conn",
369-
) as mock_validate_conn:
289+
with _mock_pool_method("_validate_conn") as mock_validate_conn:
370290
mock_validate_conn.side_effect = OSError("Some arbitrary network error")
371291
with pytest.raises(MaxRetryError) as cm:
372292
extra_params = {**extra_params, **self._retry_policy}
@@ -382,14 +302,13 @@ def test_oserror_retries(self, mock_send_telemetry, extra_params):
382302
{"use_sea": True},
383303
],
384304
)
385-
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
386-
def test_retry_max_count_not_exceeded(self, mock_send_telemetry, extra_params):
305+
def test_retry_max_count_not_exceeded(self, extra_params):
387306
"""GIVEN the max_attempts_count is 5
388307
WHEN the server sends nothing but 429 responses
389308
THEN the connector issues six request (original plus five retries)
390309
before raising an exception
391310
"""
392-
with _isolated_from_telemetry(), mocked_server_response(
311+
with mocked_server_response(
393312
status=429, headers={"Retry-After": "0"}
394313
) as mock_obj:
395314
with pytest.raises(MaxRetryError) as cm:
@@ -405,8 +324,7 @@ def test_retry_max_count_not_exceeded(self, mock_send_telemetry, extra_params):
405324
{"use_sea": True},
406325
],
407326
)
408-
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
409-
def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params):
327+
def test_retry_exponential_backoff(self, extra_params):
410328
"""GIVEN the retry policy is configured for reasonable exponential backoff
411329
WHEN the server sends nothing but 429 responses with retry-afters
412330
THEN the connector will use those retry-afters values as floor
@@ -415,7 +333,7 @@ def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params):
415333
retry_policy["_retry_delay_min"] = 1
416334

417335
time_start = time.time()
418-
with _isolated_from_telemetry(), mocked_server_response(
336+
with mocked_server_response(
419337
status=429, headers={"Retry-After": "8"}
420338
) as mock_obj:
421339
with pytest.raises(RequestError) as cm:
@@ -636,10 +554,7 @@ def test_retry_abort_close_operation_on_404(self, extra_params, caplog):
636554
{"use_sea": True},
637555
],
638556
)
639-
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
640-
def test_3xx_redirect_codes_are_not_retried(
641-
self, mock_send_telemetry, extra_params
642-
):
557+
def test_3xx_redirect_codes_are_not_retried(self, extra_params):
643558
"""GIVEN the connector is configured with a custom max_redirects
644559
WHEN the DatabricksRetryPolicy receives a 302 redirect
645560
THEN the connector does not retry since 3xx codes are not retried per policy
@@ -672,9 +587,8 @@ def test_3xx_redirect_codes_are_not_retried(
672587
{"use_sea": True},
673588
],
674589
)
675-
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
676590
def test_3xx_codes_not_retried_regardless_of_max_redirects_setting(
677-
self, mock_send_telemetry, extra_params
591+
self, extra_params
678592
):
679593
"""GIVEN the connector is configured without a custom max_redirects
680594
WHEN the DatabricksRetryPolicy receives a 302 redirect

0 commit comments

Comments
 (0)