1
1
import uuid
2
2
import pytest
3
3
from unittest .mock import patch , MagicMock
4
- import time
5
- import requests
6
- from pybreaker import CircuitBreakerError
7
- from databricks .sql .common .http import TelemetryHttpClient
8
4
9
5
from databricks .sql .telemetry .telemetry_client import (
10
6
TelemetryClient ,
11
7
NoopTelemetryClient ,
12
8
TelemetryClientFactory ,
13
9
TelemetryHelper ,
10
+ BaseTelemetryClient ,
14
11
)
15
12
from databricks .sql .telemetry .models .enums import AuthMech , AuthFlow
16
13
from databricks .sql .auth .authenticators import (
@@ -319,93 +316,3 @@ def test_connection_failure_sends_correct_telemetry_payload(
319
316
call_arguments = mock_export_failure_log .call_args
320
317
assert call_arguments [0 ][0 ] == "Exception"
321
318
assert call_arguments [0 ][1 ] == error_message
322
-
323
-
324
- class TestTelemetryHttpClient :
325
- """Tests for the TelemetryHttpClient, including retry and circuit breaker logic."""
326
-
327
- @pytest .fixture
328
- def http_client (self ):
329
- """
330
- Provides a fresh TelemetryHttpClient instance for each test,
331
- ensuring the singleton state is reset.
332
- """
333
- if TelemetryHttpClient ._instance :
334
- TelemetryHttpClient .get_instance ().close ()
335
-
336
- client = TelemetryHttpClient .get_instance ()
337
- yield client
338
-
339
- client .close ()
340
-
341
- def test_circuit_breaker_opens_after_failures (self , http_client ):
342
- """Verify the circuit opens after N consecutive failures and rejects new calls."""
343
- fail_max = 3
344
- http_client .breaker .fail_max = fail_max
345
-
346
- with patch .object (http_client .session , "post" ) as mock_post :
347
- mock_post .side_effect = requests .exceptions .RequestException ("Connection failed" )
348
-
349
- for _ in range (fail_max - 1 ):
350
- with pytest .raises (requests .exceptions .RequestException ):
351
- http_client .post ("https://test.com/telemetry" )
352
-
353
- with pytest .raises (CircuitBreakerError ):
354
- http_client .post ("https://test.com/telemetry" )
355
-
356
- assert http_client .breaker .current_state == "open"
357
- assert mock_post .call_count == fail_max
358
-
359
- with pytest .raises (CircuitBreakerError ):
360
- http_client .post ("https://test.com/telemetry" )
361
- assert mock_post .call_count == fail_max
362
-
363
- def test_circuit_breaker_closes_after_timeout_and_success (self , http_client ):
364
- """Verify the circuit moves to half-open and then closes after a successful probe."""
365
- fail_max = 2
366
- reset_timeout = 0.1
367
- http_client .breaker .fail_max = fail_max
368
- http_client .breaker .reset_timeout = reset_timeout
369
-
370
- with patch .object (http_client .session , "post" ) as mock_post :
371
- mock_post .side_effect = [
372
- requests .exceptions .RequestException ("Fail 1" ),
373
- requests .exceptions .RequestException ("Fail 2" ),
374
- MagicMock (ok = True )
375
- ]
376
-
377
- with pytest .raises (requests .exceptions .RequestException ):
378
- http_client .post ("https://test.com" )
379
- with pytest .raises (CircuitBreakerError ):
380
- http_client .post ("https://test.com" )
381
-
382
- assert http_client .breaker .current_state == "open"
383
- time .sleep (reset_timeout )
384
-
385
- http_client .post ("https://test.com" )
386
- assert http_client .breaker .current_state == "closed"
387
- assert mock_post .call_count == 3
388
-
389
- def test_circuit_breaker_reopens_if_probe_fails (self , http_client ):
390
- """Verify the circuit moves to half-open and then back to open if the probe fails."""
391
- fail_max = 2
392
- reset_timeout = 0.1
393
- http_client .breaker .fail_max = fail_max
394
- http_client .breaker .reset_timeout = reset_timeout
395
-
396
- with patch .object (http_client .session , "post" ) as mock_post :
397
- mock_post .side_effect = requests .exceptions .RequestException ("Always fails" )
398
-
399
- with pytest .raises (requests .exceptions .RequestException ):
400
- http_client .post ("https://test.com" )
401
- with pytest .raises (CircuitBreakerError ):
402
- http_client .post ("https://test.com" )
403
-
404
- assert http_client .breaker .current_state == "open"
405
- time .sleep (reset_timeout )
406
-
407
- with pytest .raises (CircuitBreakerError ):
408
- http_client .post ("https://test.com" )
409
-
410
- assert http_client .breaker .current_state == "open"
411
- assert mock_post .call_count == 3
0 commit comments