fixed failing test cases

Signed-off-by: Nikhil Suri <nikhil.suri@databricks.com>

Author: nikhilsuri-db

src/databricks/sql/telemetry/telemetry_client.py

@@ -200,13 +200,20 @@ def __init__(
 
         # Create telemetry push client based on circuit breaker enabled flag
         if client_context.telemetry_circuit_breaker_enabled:
-            # Create circuit breaker configuration with hardcoded values
-            # These values are optimized for telemetry batching and network resilience
-            circuit_breaker_config = CircuitBreakerConfig(
-                failure_threshold=0.5,  # Opens if 50%+ of calls fail
-                minimum_calls=20,  # Minimum sample size before circuit can open
-                timeout=30,  # Time window for counting failures (seconds)
-                reset_timeout=30,  # Cool-down period before retrying (seconds)
+            # Create circuit breaker configuration from client context or use defaults
+            self._circuit_breaker_config = CircuitBreakerConfig(
+                failure_threshold=getattr(
+                    client_context, "telemetry_circuit_breaker_failure_threshold", 0.5
+                ),
+                minimum_calls=getattr(
+                    client_context, "telemetry_circuit_breaker_minimum_calls", 20
+                ),
+                timeout=getattr(
+                    client_context, "telemetry_circuit_breaker_timeout", 30
+                ),
+                reset_timeout=getattr(
+                    client_context, "telemetry_circuit_breaker_reset_timeout", 30
+                ),
                 name=f"telemetry-circuit-breaker-{session_id_hex}",
             )
 
@@ -215,11 +222,12 @@ def __init__(
                 CircuitBreakerTelemetryPushClient(
                     TelemetryPushClient(self._http_client),
                     host_url,
-                    circuit_breaker_config,
+                    self._circuit_breaker_config,
                 )
             )
         else:
             # Circuit breaker disabled - use direct telemetry push client
+            self._circuit_breaker_config = None
             self._telemetry_push_client: ITelemetryPushClient = TelemetryPushClient(
                 self._http_client
             )
@@ -402,6 +410,18 @@ def close(self):
         logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
         self._flush()
 
+    def get_circuit_breaker_state(self) -> str:
+        """Get the current state of the circuit breaker."""
+        return self._telemetry_push_client.get_circuit_breaker_state()
+
+    def is_circuit_breaker_open(self) -> bool:
+        """Check if the circuit breaker is currently open."""
+        return self._telemetry_push_client.is_circuit_breaker_open()
+
+    def reset_circuit_breaker(self) -> None:
+        """Reset the circuit breaker."""
+        self._telemetry_push_client.reset_circuit_breaker()
+
 
 class TelemetryClientFactory:
     """

tests/unit/test_circuit_breaker_http_client.py

@@ -71,42 +71,25 @@ def test_initialization(self):
         assert self.client._config == self.config
         assert self.client._circuit_breaker is not None
 
-    def test_initialization_disabled(self):
-        """Test client initialization with circuit breaker disabled."""
-        config = CircuitBreakerConfig(enabled=False)
-        client = CircuitBreakerHttpClient(self.mock_delegate, self.host, config)
-        
-        assert client._config.enabled is False
 
-    def test_request_context_disabled(self):
-        """Test request context when circuit breaker is disabled."""
-        config = CircuitBreakerConfig(enabled=False)
-        client = CircuitBreakerHttpClient(self.mock_delegate, self.host, config)
-        
-        mock_response = Mock()
-        self.mock_delegate.request_context.return_value.__enter__.return_value = mock_response
-        self.mock_delegate.request_context.return_value.__exit__.return_value = None
-        
-        with client.request_context(HttpMethod.POST, "https://test.com", {}) as response:
-            assert response == mock_response
-        
-        self.mock_delegate.request_context.assert_called_once()
 
     def test_request_context_enabled_success(self):
         """Test successful request context when circuit breaker is enabled."""
         mock_response = Mock()
-        self.mock_delegate.request_context.return_value.__enter__.return_value = mock_response
-        self.mock_delegate.request_context.return_value.__exit__.return_value = None
+        mock_context = MagicMock()
+        mock_context.__enter__.return_value = mock_response
+        mock_context.__exit__.return_value = None
+        self.mock_delegate.request_context.return_value = mock_context
 
-        with client.request_context(HttpMethod.POST, "https://test.com", {}) as response:
+        with self.client.request_context(HttpMethod.POST, "https://test.com", {}) as response:
             assert response == mock_response
 
         self.mock_delegate.request_context.assert_called_once()
 
     def test_request_context_enabled_circuit_breaker_error(self):
         """Test request context when circuit breaker is open."""
         # Mock circuit breaker to raise CircuitBreakerError
-        with patch.object(self.client._circuit_breaker, '__enter__', side_effect=CircuitBreakerError("Circuit is open")):
+        with patch.object(self.client._circuit_breaker, 'call', side_effect=CircuitBreakerError("Circuit is open")):
             with pytest.raises(CircuitBreakerError):
                 with self.client.request_context(HttpMethod.POST, "https://test.com", {}):
                     pass
@@ -120,18 +103,6 @@ def test_request_context_enabled_other_error(self):
             with self.client.request_context(HttpMethod.POST, "https://test.com", {}):
                 pass
 
-    def test_request_disabled(self):
-        """Test request method when circuit breaker is disabled."""
-        config = CircuitBreakerConfig(enabled=False)
-        client = CircuitBreakerHttpClient(self.mock_delegate, self.host, config)
-        
-        mock_response = Mock()
-        self.mock_delegate.request.return_value = mock_response
-        
-        response = client.request(HttpMethod.POST, "https://test.com", {})
-        
-        assert response == mock_response
-        self.mock_delegate.request.assert_called_once()
 
     def test_request_enabled_success(self):
         """Test successful request when circuit breaker is enabled."""
@@ -146,7 +117,7 @@ def test_request_enabled_success(self):
     def test_request_enabled_circuit_breaker_error(self):
         """Test request when circuit breaker is open."""
         # Mock circuit breaker to raise CircuitBreakerError
-        with patch.object(self.client._circuit_breaker, '__enter__', side_effect=CircuitBreakerError("Circuit is open")):
+        with patch.object(self.client._circuit_breaker, 'call', side_effect=CircuitBreakerError("Circuit is open")):
             with pytest.raises(CircuitBreakerError):
                 self.client.request(HttpMethod.POST, "https://test.com", {})
 
@@ -160,15 +131,15 @@ def test_request_enabled_other_error(self):
 
     def test_get_circuit_breaker_state(self):
         """Test getting circuit breaker state."""
-        with patch.object(self.client._circuit_breaker, 'current_state', 'open'):
+        with patch('databricks.sql.telemetry.telemetry_push_client.CircuitBreakerManager.get_circuit_breaker_state', return_value='open'):
             state = self.client.get_circuit_breaker_state()
             assert state == 'open'
 
     def test_reset_circuit_breaker(self):
         """Test resetting circuit breaker."""
-        with patch.object(self.client._circuit_breaker, 'reset') as mock_reset:
+        with patch('databricks.sql.telemetry.telemetry_push_client.CircuitBreakerManager.reset_circuit_breaker') as mock_reset:
             self.client.reset_circuit_breaker()
-            mock_reset.assert_called_once()
+            mock_reset.assert_called_once_with(self.client._host)
 
     def test_is_circuit_breaker_open(self):
         """Test checking if circuit breaker is open."""
@@ -180,42 +151,38 @@ def test_is_circuit_breaker_open(self):
 
     def test_is_circuit_breaker_enabled(self):
         """Test checking if circuit breaker is enabled."""
-        assert self.client.is_circuit_breaker_enabled() is True
-        
-        config = CircuitBreakerConfig(enabled=False)
-        client = CircuitBreakerHttpClient(self.mock_delegate, self.host, config)
-        assert client.is_circuit_breaker_enabled() is False
+        assert self.client._circuit_breaker is not None
 
     def test_circuit_breaker_state_logging(self):
         """Test that circuit breaker state changes are logged."""
-        with patch('databricks.sql.telemetry.circuit_breaker_http_client.logger') as mock_logger:
-            with patch.object(self.client._circuit_breaker, '__enter__', side_effect=CircuitBreakerError("Circuit is open")):
+        with patch('databricks.sql.telemetry.telemetry_push_client.logger') as mock_logger:
+            with patch.object(self.client._circuit_breaker, 'call', side_effect=CircuitBreakerError("Circuit is open")):
                 with pytest.raises(CircuitBreakerError):
                     self.client.request(HttpMethod.POST, "https://test.com", {})
-                
-                # Check that warning was logged
-                mock_logger.warning.assert_called()
-                warning_call = mock_logger.warning.call_args[0][0]
-                assert "Circuit breaker is open" in warning_call
-                assert self.host in warning_call
+            
+            # Check that warning was logged
+            mock_logger.warning.assert_called()
+            warning_call = mock_logger.warning.call_args[0]
+            assert "Circuit breaker is open" in warning_call[0]
+            assert self.host in warning_call[1]
 
     def test_other_error_logging(self):
         """Test that other errors are logged appropriately."""
-        with patch('databricks.sql.telemetry.circuit_breaker_http_client.logger') as mock_logger:
+        with patch('databricks.sql.telemetry.telemetry_push_client.logger') as mock_logger:
             self.mock_delegate.request.side_effect = ValueError("Network error")
 
             with pytest.raises(ValueError):
                 self.client.request(HttpMethod.POST, "https://test.com", {})
 
             # Check that debug was logged
             mock_logger.debug.assert_called()
-            debug_call = mock_logger.debug.call_args[0][0]
-            assert "Telemetry request failed" in debug_call
-            assert self.host in debug_call
+            debug_call = mock_logger.debug.call_args[0]
+            assert "Telemetry request failed" in debug_call[0]
+            assert self.host in debug_call[1]
 
 
-class TestCircuitBreakerHttpClientIntegration:
-    """Integration tests for CircuitBreakerHttpClient."""
+class TestCircuitBreakerTelemetryPushClientIntegration:
+    """Integration tests for CircuitBreakerTelemetryPushClient."""
 
     def setup_method(self):
         """Set up test fixtures."""
@@ -224,42 +191,59 @@ def setup_method(self):
 
     def test_circuit_breaker_opens_after_failures(self):
         """Test that circuit breaker opens after repeated failures."""
+        from databricks.sql.telemetry.circuit_breaker_manager import CircuitBreakerManager
+        
+        # Clear any existing state
+        CircuitBreakerManager.clear_all_circuit_breakers()
+        
         config = CircuitBreakerConfig(
             failure_threshold=0.1,  # 10% failure rate
             minimum_calls=2,        # Only 2 calls needed
             reset_timeout=1         # 1 second reset timeout
         )
-        client = CircuitBreakerHttpClient(self.mock_delegate, self.host, config)
+        
+        # Initialize the manager
+        CircuitBreakerManager.initialize(config)
+        
+        client = CircuitBreakerTelemetryPushClient(self.mock_delegate, self.host, config)
 
         # Simulate failures
         self.mock_delegate.request.side_effect = Exception("Network error")
 
-        # First few calls should fail with the original exception
-        for _ in range(2):
-            with pytest.raises(Exception, match="Network error"):
-                client.request(HttpMethod.POST, "https://test.com", {})
+        # First call should fail with the original exception
+        with pytest.raises(Exception, match="Network error"):
+            client.request(HttpMethod.POST, "https://test.com", {})
 
-        # After enough failures, circuit breaker should open
+        # Second call should open the circuit breaker and raise CircuitBreakerError
         with pytest.raises(CircuitBreakerError):
             client.request(HttpMethod.POST, "https://test.com", {})
 
     def test_circuit_breaker_recovers_after_success(self):
         """Test that circuit breaker recovers after successful calls."""
+        from databricks.sql.telemetry.circuit_breaker_manager import CircuitBreakerManager
+        
+        # Clear any existing state
+        CircuitBreakerManager.clear_all_circuit_breakers()
+        
         config = CircuitBreakerConfig(
             failure_threshold=0.1,
             minimum_calls=2,
             reset_timeout=1
         )
-        client = CircuitBreakerHttpClient(self.mock_delegate, self.host, config)
+        
+        # Initialize the manager
+        CircuitBreakerManager.initialize(config)
+        
+        client = CircuitBreakerTelemetryPushClient(self.mock_delegate, self.host, config)
 
         # Simulate failures first
         self.mock_delegate.request.side_effect = Exception("Network error")
 
-        for _ in range(2):
-            with pytest.raises(Exception):
-                client.request(HttpMethod.POST, "https://test.com", {})
+        # First call should fail with the original exception
+        with pytest.raises(Exception):
+            client.request(HttpMethod.POST, "https://test.com", {})
 
-        # Circuit breaker should be open now
+        # Second call should open the circuit breaker
         with pytest.raises(CircuitBreakerError):
             client.request(HttpMethod.POST, "https://test.com", {})
 

tests/unit/test_circuit_breaker_manager.py

@@ -78,7 +78,7 @@ def test_get_circuit_breaker_not_initialized(self):
 
         # Should return a no-op circuit breaker
         assert breaker.name == "noop-circuit-breaker"
-        assert breaker.failure_threshold == 1.0
+        assert breaker.fail_max == 1000000  # Very high threshold for no-op
 
     def test_get_circuit_breaker_enabled(self):
         """Test getting circuit breaker when enabled."""