work

databricks/sdk/config.py

@@ -373,7 +373,8 @@ def load_azure_tenant_id(self):
         logger.debug(f'Loading tenant ID from {login_url}')
         resp = requests.get(login_url, allow_redirects=False)
         if resp.status_code // 100 != 3:
-            logger.debug(f'Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status_code}')
+            logger.debug(
+                f'Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status_code}')
             return
         entra_id_endpoint = resp.headers.get('Location')
         if entra_id_endpoint is None:

databricks/sdk/credentials_provider.py

@@ -233,11 +233,11 @@ def _ensure_host_present(cfg: 'Config', token_source_for: Callable[[str], TokenS
     cfg.host = f"https://{resp.json()['properties']['workspaceUrl']}"
 
 
-@oauth_credentials_strategy('azure-client-secret',
-                            ['is_azure', 'azure_client_id', 'azure_client_secret'])
+@oauth_credentials_strategy('azure-client-secret', ['is_azure', 'azure_client_id', 'azure_client_secret'])
 def azure_service_principal(cfg: 'Config') -> CredentialsProvider:
     """ Adds refreshed Azure Active Directory (AAD) Service Principal OAuth tokens
     to every request, while automatically resolving different Azure environment endpoints. """
+
     def token_source_for(resource: str) -> TokenSource:
         aad_endpoint = cfg.arm_environment.active_directory_endpoint
         return ClientCredentials(client_id=cfg.azure_client_id,
@@ -467,7 +467,9 @@ def is_human_user(self) -> bool:
     def for_resource(cfg: 'Config', resource: str) -> 'AzureCliTokenSource':
         subscription = AzureCliTokenSource.get_subscription(cfg)
         if subscription is not None:
-            token_source = AzureCliTokenSource(resource, subscription=subscription, tenant=cfg.azure_tenant_id)
+            token_source = AzureCliTokenSource(resource,
+                                               subscription=subscription,
+                                               tenant=cfg.azure_tenant_id)
             try:
                 # This will fail if the user has access to the workspace, but not to the subscription
                 # itself.

tests/conftest.py

@@ -77,3 +77,16 @@ def set_az_path(monkeypatch):
         monkeypatch.setenv('COMSPEC', 'C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe')
     else:
         monkeypatch.setenv('PATH', __tests__ + "/testdata:/bin")
+
+
+@pytest.fixture
+def mock_tenant(requests_mock):
+
+    def stub_tenant_request(host, tenant_id="test-tenant-id"):
+        mock = requests_mock.get(
+            f'https://{host}/aad/auth',
+            status_code=302,
+            headers={'Location': f'https://login.microsoftonline.com/{tenant_id}/oauth2/authorize'})
+        return mock
+
+    return stub_tenant_request

tests/test_auth.py

@@ -193,9 +193,10 @@ def test_config_azure_pat():
     assert cfg.is_azure
 
 
-def test_config_azure_cli_host(monkeypatch):
+def test_config_azure_cli_host(monkeypatch, mock_tenant):
     set_home(monkeypatch, '/testdata/azure')
     set_az_path(monkeypatch)
+    mock_tenant('adb-123.4.azuredatabricks.net')
     cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws')
 
     assert cfg.auth_type == 'azure-cli'
@@ -229,20 +230,22 @@ def test_config_azure_cli_host_pat_conflict_with_config_file_present_without_def
     cfg = Config(token='x', azure_workspace_resource_id='/sub/rg/ws')
 
 
-def test_config_azure_cli_host_and_resource_id(monkeypatch):
+def test_config_azure_cli_host_and_resource_id(monkeypatch, mock_tenant):
     set_home(monkeypatch, '/testdata')
     set_az_path(monkeypatch)
+    mock_tenant('adb-123.4.azuredatabricks.net')
     cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws')
 
     assert cfg.auth_type == 'azure-cli'
     assert cfg.host == 'https://adb-123.4.azuredatabricks.net'
     assert cfg.is_azure
 
 
-def test_config_azure_cli_host_and_resource_i_d_configuration_precedence(monkeypatch):
+def test_config_azure_cli_host_and_resource_i_d_configuration_precedence(monkeypatch, mock_tenant):
     monkeypatch.setenv('DATABRICKS_CONFIG_PROFILE', 'justhost')
     set_home(monkeypatch, '/testdata/azure')
     set_az_path(monkeypatch)
+    mock_tenant('adb-123.4.azuredatabricks.net')
     cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws')
 
     assert cfg.auth_type == 'azure-cli'

tests/test_auth_manual_tests.py

@@ -3,9 +3,10 @@
 from .conftest import set_az_path, set_home
 
 
-def test_azure_cli_workspace_header_present(monkeypatch):
+def test_azure_cli_workspace_header_present(monkeypatch, mock_tenant):
     set_home(monkeypatch, '/testdata/azure')
     set_az_path(monkeypatch)
+    mock_tenant('adb-123.4.azuredatabricks.net')
     resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
     cfg = Config(auth_type='azure-cli',
                  host='https://adb-123.4.azuredatabricks.net',
@@ -14,19 +15,21 @@ def test_azure_cli_workspace_header_present(monkeypatch):
     assert cfg.authenticate()['X-Databricks-Azure-Workspace-Resource-Id'] == resource_id
 
 
-def test_azure_cli_user_with_management_access(monkeypatch):
+def test_azure_cli_user_with_management_access(monkeypatch, mock_tenant):
     set_home(monkeypatch, '/testdata/azure')
     set_az_path(monkeypatch)
+    mock_tenant('adb-123.4.azuredatabricks.net')
     resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
     cfg = Config(auth_type='azure-cli',
                  host='https://adb-123.4.azuredatabricks.net',
                  azure_workspace_resource_id=resource_id)
     assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate()
 
 
-def test_azure_cli_user_no_management_access(monkeypatch):
+def test_azure_cli_user_no_management_access(monkeypatch, mock_tenant):
     set_home(monkeypatch, '/testdata/azure')
     set_az_path(monkeypatch)
+    mock_tenant('adb-123.4.azuredatabricks.net')
     monkeypatch.setenv('FAIL_IF', 'https://management.core.windows.net/')
     resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
     cfg = Config(auth_type='azure-cli',
@@ -35,9 +38,10 @@ def test_azure_cli_user_no_management_access(monkeypatch):
     assert 'X-Databricks-Azure-SP-Management-Token' not in cfg.authenticate()
 
 
-def test_azure_cli_fallback(monkeypatch):
+def test_azure_cli_fallback(monkeypatch, mock_tenant):
     set_home(monkeypatch, '/testdata/azure')
     set_az_path(monkeypatch)
+    mock_tenant('adb-123.4.azuredatabricks.net')
     monkeypatch.setenv('FAIL_IF', 'subscription')
     resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
     cfg = Config(auth_type='azure-cli',
@@ -46,9 +50,10 @@ def test_azure_cli_fallback(monkeypatch):
     assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate()
 
 
-def test_azure_cli_with_warning_on_stderr(monkeypatch):
+def test_azure_cli_with_warning_on_stderr(monkeypatch, mock_tenant):
     set_home(monkeypatch, '/testdata/azure')
     set_az_path(monkeypatch)
+    mock_tenant('adb-123.4.azuredatabricks.net')
     monkeypatch.setenv('WARN', 'this is a warning')
     resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
     cfg = Config(auth_type='azure-cli',

tests/test_config.py

@@ -1,3 +1,4 @@
+import os
 import platform
 
 import pytest
@@ -8,10 +9,9 @@
 
 from .conftest import noop_credentials
 
-import os
-
 __tests__ = os.path.dirname(__file__)
 
+
 def test_config_supports_legacy_credentials_provider():
     c = Config(credentials_provider=noop_credentials, product='foo', product_version='1.2.3')
     c2 = c.copy()
@@ -97,15 +97,20 @@ def test_load_azure_tenant_id_no_location_header(requests_mock, monkeypatch):
 
 def test_load_azure_tenant_id_unparsable_location_header(requests_mock, monkeypatch):
     monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
-    mock = requests_mock.get('https://abc123.azuredatabricks.net/aad/auth', status_code=302, headers={'Location': 'https://unexpected-location'})
+    mock = requests_mock.get('https://abc123.azuredatabricks.net/aad/auth',
+                             status_code=302,
+                             headers={'Location': 'https://unexpected-location'})
     cfg = Config(host="https://abc123.azuredatabricks.net")
     assert cfg.azure_tenant_id is None
     assert mock.called_once
 
 
 def test_load_azure_tenant_id_happy_path(requests_mock, monkeypatch):
     monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
-    mock = requests_mock.get('https://abc123.azuredatabricks.net/aad/auth', status_code=302, headers={'Location': 'https://login.microsoftonline.com/tenant-id/oauth2/authorize'})
+    mock = requests_mock.get(
+        'https://abc123.azuredatabricks.net/aad/auth',
+        status_code=302,
+        headers={'Location': 'https://login.microsoftonline.com/tenant-id/oauth2/authorize'})
     cfg = Config(host="https://abc123.azuredatabricks.net")
     assert cfg.azure_tenant_id == 'tenant-id'
     assert mock.called_once