From 6c014ca51f56a52482df6068cc4477dac417eaa5 Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Tue, 6 Aug 2024 13:14:06 +0200
Subject: [PATCH] Decrease runtime of recursive workspace listing test

---
 tests/integration/conftest.py       |  8 ++++++++
 tests/integration/test_workspace.py | 16 +++++++++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index 221cc7d49..e16335bf3 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -104,6 +104,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
diff --git a/tests/integration/test_workspace.py b/tests/integration/test_workspace.py
index afe77c427..4adbee773 100644
--- a/tests/integration/test_workspace.py
+++ b/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):