[Fix] Fix test_get_workspace_client and test_runtime_auth_from_jobs (#…

…719)

## Changes
This PR fixes the current failing integration tests for the Python SDK,
unblocking their release.

There are two issues:
1. get_workspace_client fails in our integration tests because we call
it with a workspace that is not UC-enabled. Because tests are
authenticated as service principals, and it isn't possible to add
account-level service principals to non-UC workspaces, this call fails.
I address this by running this test against a UC-enabled workspace.
2. test_runtime_auth_from_jobs fails because a new LTS DBR version was
released (15.4) that doesn't support DBFS library installations. To
address this, I have created two tests:
test_runtime_auth_from_jobs_dbfs, which tests native auth using the SDK
installed from DBFS up to LTS 14.3, and
test_runtime_auth_from_jobs_volumes, which does the same with the SDK
installed from a volume.

## Tests
All integration tests passed (retriggered the GCP integration test
locally after adding single user data security mode).

- [ ] `make test` run locally
- [ ] `make fmt` applied
- [ ] relevant integration tests applied

tests/integration/conftest.py

@@ -59,6 +59,18 @@ def a(env_or_skip) -> AccountClient:
     return account_client
 
 
+@pytest.fixture(scope='session')
+def ucacct(env_or_skip) -> AccountClient:
+    _load_debug_env_if_runs_from_ide('ucacct')
+    env_or_skip("CLOUD_ENV")
+    account_client = AccountClient()
+    if not account_client.config.is_account_client:
+        pytest.skip("not Databricks Account client")
+    if 'TEST_METASTORE_ID' not in os.environ:
+        pytest.skip("not in Unity Catalog Workspace test env")
+    return account_client
+
+
 @pytest.fixture(scope='session')
 def w(env_or_skip) -> WorkspaceClient:
     _load_debug_env_if_runs_from_ide('workspace')

tests/integration/test_auth.py

@@ -5,14 +5,15 @@
 import shutil
 import subprocess
 import sys
+import typing
 import urllib.parse
 from functools import partial
 from pathlib import Path
 
 import pytest
 
 from databricks.sdk.service.compute import (ClusterSpec, DataSecurityMode,
-                                            Library, ResultType)
+                                            Library, ResultType, SparkVersion)
 from databricks.sdk.service.jobs import NotebookTask, Task, ViewType
 from databricks.sdk.service.workspace import ImportFormat
 
@@ -84,19 +85,41 @@ def test_runtime_auth_from_interactive_on_uc(ucws, fresh_wheel_file, env_or_skip
         ucws.clusters.permanent_delete(interactive_cluster.cluster_id)
 
 
-def test_runtime_auth_from_jobs(w, fresh_wheel_file, env_or_skip, random):
-    instance_pool_id = env_or_skip('TEST_INSTANCE_POOL_ID')
-
+def _get_lts_versions(w) -> typing.List[SparkVersion]:
     v = w.clusters.spark_versions()
     lts_runtimes = [
         x for x in v.versions
         if 'LTS' in x.name and '-ml' not in x.key and '-photon' not in x.key and '-aarch64' not in x.key
     ]
+    return lts_runtimes
+
+
+def test_runtime_auth_from_jobs_volumes(ucws, fresh_wheel_file, env_or_skip, random, volume):
+    dbr_versions = [v for v in _get_lts_versions(ucws) if int(v.key.split('.')[0]) >= 15]
+
+    volume_wheel = f'{volume}/tmp/wheels/{random(10)}/{fresh_wheel_file.name}'
+    with fresh_wheel_file.open('rb') as f:
+        ucws.files.upload(volume_wheel, f)
+
+    lib = Library(whl=volume_wheel)
+    return _test_runtime_auth_from_jobs_inner(ucws, env_or_skip, random, dbr_versions, lib)
+
+
+def test_runtime_auth_from_jobs_dbfs(w, fresh_wheel_file, env_or_skip, random):
+    # Library installation from DBFS is not supported past DBR 14.3
+    dbr_versions = [v for v in _get_lts_versions(w) if int(v.key.split('.')[0]) < 15]
 
     dbfs_wheel = f'/tmp/wheels/{random(10)}/{fresh_wheel_file.name}'
     with fresh_wheel_file.open('rb') as f:
         w.dbfs.upload(dbfs_wheel, f)
 
+    lib = Library(whl=f'dbfs:{dbfs_wheel}')
+    return _test_runtime_auth_from_jobs_inner(w, env_or_skip, random, dbr_versions, lib)
+
+
+def _test_runtime_auth_from_jobs_inner(w, env_or_skip, random, dbr_versions, library):
+    instance_pool_id = env_or_skip('TEST_INSTANCE_POOL_ID')
+
     my_name = w.current_user.me().user_name
     notebook_path = f'/Users/{my_name}/notebook-native-auth'
     notebook_content = io.BytesIO(b'''
@@ -109,16 +132,20 @@ def test_runtime_auth_from_jobs(w, fresh_wheel_file, env_or_skip, random):
     w.workspace.upload(notebook_path, notebook_content, language=Language.PYTHON, overwrite=True)
 
     tasks = []
-    for v in lts_runtimes:
+    for v in dbr_versions:
         t = Task(task_key=f'test_{v.key.replace(".", "_")}',
                  notebook_task=NotebookTask(notebook_path=notebook_path),
-                 new_cluster=ClusterSpec(spark_version=v.key,
-                                         num_workers=1,
-                                         instance_pool_id=instance_pool_id),
-                 libraries=[Library(whl=f'dbfs:{dbfs_wheel}')])
+                 new_cluster=ClusterSpec(
+                     spark_version=v.key,
+                     num_workers=1,
+                     instance_pool_id=instance_pool_id,
+                     # GCP uses "custom" data security mode by default, which does not support UC.
+                     data_security_mode=DataSecurityMode.SINGLE_USER),
+                 libraries=[library])
         tasks.append(t)
 
-    run = w.jobs.submit(run_name=f'Runtime Native Auth {random(10)}', tasks=tasks).result()
+    waiter = w.jobs.submit(run_name=f'Runtime Native Auth {random(10)}', tasks=tasks)
+    run = waiter.result()
     for task_key, output in _task_outputs(w, run).items():
         assert my_name in output, f'{task_key} does not work with notebook native auth'
 

tests/integration/test_client.py

@@ -1,10 +1,11 @@
 import pytest
 
 
-def test_get_workspace_client(a, env_or_skip):
+def test_get_workspace_client(ucacct, env_or_skip):
+    # Need to switch to ucacct
     workspace_id = env_or_skip("TEST_WORKSPACE_ID")
-    ws = a.workspaces.get(workspace_id)
-    w = a.get_workspace_client(ws)
+    ws = ucacct.workspaces.get(workspace_id)
+    w = ucacct.get_workspace_client(ws)
     assert w.current_user.me().active