Merge branch 'main' into main

Author: pivoshenko

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')
@@ -104,6 +116,14 @@ def volume(ucws, schema):
     ucws.volumes.delete(volume.full_name)
 
 
+@pytest.fixture()
+def workspace_dir(w, random):
+    directory = f'/Users/{w.current_user.me().user_name}/dir-{random(12)}'
+    w.workspace.mkdirs(directory)
+    yield directory
+    w.workspace.delete(directory, recursive=True)
+
+
 def _load_debug_env_if_runs_from_ide(key) -> bool:
     if not _is_in_debug():
         return False

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
 
 

tests/integration/test_workspace.py

@@ -3,11 +3,21 @@
 from databricks.sdk.service.workspace import ImportFormat, Language
 
 
-def test_workspace_recursive_list(w, random):
+def test_workspace_recursive_list(w, workspace_dir, random):
+    # create a file in the directory
+    file = f'{workspace_dir}/file-{random(12)}.py'
+    w.workspace.upload(file, io.BytesIO(b'print(1)'))
+    # create a subdirectory
+    subdirectory = f'{workspace_dir}/subdir-{random(12)}'
+    w.workspace.mkdirs(subdirectory)
+    # create a file in the subdirectory
+    subfile = f'{subdirectory}/subfile-{random(12)}.py'
+    w.workspace.upload(subfile, io.BytesIO(b'print(2)'))
+    # list the directory recursively
     names = []
-    for i in w.workspace.list(f'/Users/{w.current_user.me().user_name}', recursive=True):
+    for i in w.workspace.list(workspace_dir, recursive=True):
         names.append(i.path)
-    assert len(names) > 0
+    assert len(names) == 2
 
 
 def test_workspace_upload_download_notebooks(w, random):