2424AzureOAuthEndpointCollection ,
2525)
2626from databricks .sql .auth .authenticators import CredentialsProvider ,HeaderFactory
27+ from databricks .sql .common .http import DatabricksHttpClient
2728from databricks .sql .experimental .oauth_persistence import OAuthPersistenceCache
2829
2930
@@ -248,8 +249,10 @@ def test_no_token_refresh__when_token_is_not_expired(
248249assert mock_get_token .call_count == 1
249250
250251def test_get_token_success (self ,token_source ,http_response ):
251- with patch .object (token_source ._http_client ,"execute" )as mock_execute :
252- mock_execute .return_value = http_response (200 )
252+ databricks_http_client = DatabricksHttpClient .get_instance ()
253+ with patch .object (
254+ databricks_http_client .session ,"request" ,return_value = http_response (200 )
255+ )as mock_request :
253256token = token_source .get_token ()
254257
255258# Assert
@@ -259,8 +262,10 @@ def test_get_token_success(self, token_source, http_response):
259262assert token .refresh_token is None
260263
261264def test_get_token_failure (self ,token_source ,http_response ):
262- with patch .object (token_source ._http_client ,"execute" )as mock_execute :
263- mock_execute .return_value = http_response (400 )
265+ databricks_http_client = DatabricksHttpClient .get_instance ()
266+ with patch .object (
267+ databricks_http_client .session ,"request" ,return_value = http_response (400 )
268+ )as mock_request :
264269with pytest .raises (Exception )as e :
265270token_source .get_token ()
266271assert "Failed to get token: 400" in str (e .value )