Skip to content

test: isolate retry HTTP mocks from background telemetry #1055

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vuanhphung merged 3 commits into from
Nov 14, 2026
Merged

test: isolate retry HTTP mocks from background telemetry #1055

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9954f0c
test: isolate retry HTTP mocks from background telemetry
vuanhphung Nov 13, 2026
e0dd510
test: keep mocked connections out of shared HTTP pools
vuanhphung Nov 13, 2026
415bb24
test: isolate pool cleanup on transport failures
vuanhphung Nov 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
166 changes: 40 additions & 126 deletions tests/e2e/common/retry_test_mixins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
from contextlib import contextmanager
from contextlib import contextmanager, ExitStack
import threading
import time
from typing import Optional, List
from unittest.mock import MagicMock, PropertyMock, patch
import io

import pytest
from urllib3.connectionpool import HTTPSConnectionPool
from urllib3.exceptions import MaxRetryError

from databricks.sql.auth.retry import DatabricksRetryPolicy
Expand All @@ -15,97 +17,35 @@
SessionAlreadyClosedError,
UnsafeToRetryError,
)
from databricks.sql.telemetry.telemetry_client import (
NoopTelemetryClient,
TelemetryClient,
TelemetryClientFactory,
_TelemetryClientHolder,
)


@contextmanager
def _isolated_from_telemetry():
# Tests that mock urllib3 globally (via _get_conn / _validate_conn) also
# intercept background telemetry pushes from the shared
# TelemetryClientFactory executor — inflating mock.call_count and breaking
# assertions like `call_count == 6`. Three layers of defence:
#
# 1. Drain TelemetryClientFactory and override initialize_telemetry_client
# so new connections install NoopTelemetryClient (which submits nothing).
# 2. Patch TelemetryClient._send_telemetry to a no-op as a backstop — covers
# any real TelemetryClient instance that slips in (e.g. a stale module-
# global, a code path that bypasses initialize_telemetry_client, or
# anything created before this context entered).
# 3. Patch TelemetryClient._export_event to a no-op so even if a real
# client receives an event, the event never reaches the queue and the
# flush logic never fires.
#
# Without layer 2/3 we have observed `/telemetry-ext` requests showing up
# in merge_group runs (concurrent CI load on the warehouse stresses paths
# that single-test runs don't hit), inflating retry-test counts and
# producing flaky AssertionErrors.
with TelemetryClientFactory._lock:
saved_clients = TelemetryClientFactory._clients
saved_executor = TelemetryClientFactory._executor
saved_initialized = TelemetryClientFactory._initialized
for holder in list(saved_clients.values()):
try:
holder.client.close()
except Exception:
pass
if saved_executor is not None:
try:
TelemetryClientFactory._stop_flush_thread()
saved_executor.shutdown(wait=True)
except Exception:
pass
TelemetryClientFactory._clients = {}
TelemetryClientFactory._executor = None
TelemetryClientFactory._initialized = False

def _install_noop(*args, host_url=None, **kwargs):
host_url = TelemetryClientFactory.getHostUrlSafely(host_url)
with TelemetryClientFactory._lock:
TelemetryClientFactory._clients[host_url] = _TelemetryClientHolder(
NoopTelemetryClient()
def _mock_pool_method(method_name):
"""Mock the calling test thread; worker threads use the original pool methods."""
original = getattr(HTTPSConnectionPool, method_name)
test_thread = threading.current_thread()
mock_method = MagicMock()

def dispatch(pool, *args, **kwargs):
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Resolved
if threading.current_thread() is test_thread:
return mock_method(*args, **kwargs)
return original(pool, *args, **kwargs)

with ExitStack() as stack:
stack.enter_context(patch.object(HTTPSConnectionPool, method_name, dispatch))
if method_name == "_get_conn":
mock_method.return_value.close.return_value = None
original_put = HTTPSConnectionPool._put_conn

def put_connection(pool, connection):
# Fake acquisition never removed a slot from the real queue.
if threading.current_thread() is not test_thread:
return original_put(pool, connection)
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Resolved

stack.enter_context(
patch.object(HTTPSConnectionPool, "_put_conn", put_connection)
)

def _noop_send(self, events):
return None

def _noop_export(self, event):
return None

try:
def _noop_send_with_unified_client(self, *args, **kwargs):
return None

with patch.object(
TelemetryClientFactory,
"initialize_telemetry_client",
staticmethod(_install_noop),
), patch.object(
TelemetryClient, "_send_telemetry", _noop_send
), patch.object(
TelemetryClient, "_export_event", _noop_export
), patch.object(
# Backstop layer 4: no-op the actual HTTP boundary. _send_telemetry
# only suppresses NEW submissions; a future already submitted (e.g. by
# the holder.client.close() flush above, which runs before these
# patches apply) still calls _send_with_unified_client on a background
# thread and would hit the network inside the test's urllib3 mock.
# Patching the HTTP method itself closes that already-submitted-future
# gap for any caller of this helper.
TelemetryClient,
"_send_with_unified_client",
_noop_send_with_unified_client,
):
yield
finally:
with TelemetryClientFactory._lock:
TelemetryClientFactory._clients = saved_clients
TelemetryClientFactory._executor = saved_executor
TelemetryClientFactory._initialized = saved_initialized
yield mock_method


class Client429ResponseMixin:
Expand Down Expand Up @@ -260,7 +200,7 @@ def mocked_server_response(

mock_response = SimpleHttpResponse(status, headers, redirect_location)

with patch("urllib3.connectionpool.HTTPSConnectionPool._get_conn") as getconn_mock:
with _mock_pool_method("_get_conn") as getconn_mock:
getconn_mock.return_value.getresponse.return_value = mock_response
try:
yield getconn_mock
Expand Down Expand Up @@ -288,7 +228,7 @@ def mock_sequential_server_responses(responses: List[dict]):
for resp in responses
]

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

Expand All @@ -331,10 +257,7 @@ class PySQLRetryTestsMixin:
{"use_sea": True},
],
)
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
def test_retry_urllib3_settings_are_honored(
self, mock_send_telemetry, extra_params
):
def test_retry_urllib3_settings_are_honored(self, extra_params):
"""Databricks overrides some of urllib3's configuration. This tests confirms that what configuration
we DON'T override is preserved in urllib3's internals
"""
Expand All @@ -361,12 +284,9 @@ def test_retry_urllib3_settings_are_honored(
{"use_sea": True},
],
)
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
def test_oserror_retries(self, mock_send_telemetry, extra_params):
def test_oserror_retries(self, extra_params):
"""If a network error occurs during make_request, the request is retried according to policy"""
with _isolated_from_telemetry(), patch(
"urllib3.connectionpool.HTTPSConnectionPool._validate_conn",
) as mock_validate_conn:
with _mock_pool_method("_validate_conn") as mock_validate_conn:
mock_validate_conn.side_effect = OSError("Some arbitrary network error")
with pytest.raises(MaxRetryError) as cm:
extra_params = {**extra_params, **self._retry_policy}
Expand All @@ -382,14 +302,13 @@ def test_oserror_retries(self, mock_send_telemetry, extra_params):
{"use_sea": True},
],
)
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
def test_retry_max_count_not_exceeded(self, mock_send_telemetry, extra_params):
def test_retry_max_count_not_exceeded(self, extra_params):
"""GIVEN the max_attempts_count is 5
WHEN the server sends nothing but 429 responses
THEN the connector issues six request (original plus five retries)
before raising an exception
"""
with _isolated_from_telemetry(), mocked_server_response(
with mocked_server_response(
status=429, headers={"Retry-After": "0"}
) as mock_obj:
with pytest.raises(MaxRetryError) as cm:
Expand All @@ -405,8 +324,7 @@ def test_retry_max_count_not_exceeded(self, mock_send_telemetry, extra_params):
{"use_sea": True},
],
)
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params):
def test_retry_exponential_backoff(self, extra_params):
"""GIVEN the retry policy is configured for reasonable exponential backoff
WHEN the server sends nothing but 429 responses with retry-afters
THEN the connector will use those retry-afters values as floor
Expand All @@ -415,7 +333,7 @@ def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params):
retry_policy["_retry_delay_min"] = 1

time_start = time.time()
with _isolated_from_telemetry(), mocked_server_response(
with mocked_server_response(
status=429, headers={"Retry-After": "8"}
) as mock_obj:
with pytest.raises(RequestError) as cm:
Expand Down Expand Up @@ -636,10 +554,7 @@ def test_retry_abort_close_operation_on_404(self, extra_params, caplog):
{"use_sea": True},
],
)
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
def test_3xx_redirect_codes_are_not_retried(
self, mock_send_telemetry, extra_params
):
def test_3xx_redirect_codes_are_not_retried(self, extra_params):
"""GIVEN the connector is configured with a custom max_redirects
WHEN the DatabricksRetryPolicy receives a 302 redirect
THEN the connector does not retry since 3xx codes are not retried per policy
Expand Down Expand Up @@ -672,9 +587,8 @@ def test_3xx_redirect_codes_are_not_retried(
{"use_sea": True},
],
)
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
def test_3xx_codes_not_retried_regardless_of_max_redirects_setting(
self, mock_send_telemetry, extra_params
self, extra_params
):
"""GIVEN the connector is configured without a custom max_redirects
WHEN the DatabricksRetryPolicy receives a 302 redirect
Expand Down
Loading
Loading