Create workspace file for SparkPythonTask

Author: JCZuurmond

src/databricks/labs/pytester/fixtures/compute.py

@@ -162,7 +162,14 @@ def create(*, instance_pool_name=None, node_type_id=None, **kwargs) -> CreateIns
 
 
 @fixture
-def make_job(ws, make_random, make_notebook, log_workspace_link, watchdog_remove_after) -> Generator[Job, None, None]:
+def make_job(
+    ws,
+    make_random,
+    make_notebook,
+    make_workspace_file,
+    log_workspace_link,
+    watchdog_remove_after,
+) -> Generator[Job, None, None]:
     """
     Create a Databricks job and clean it up after the test. Returns a function to create jobs, that returns
     a `databricks.sdk.service.jobs.Job` instance.
@@ -212,7 +219,6 @@ def create(
             raise ValueError(
                 "The `tasks` parameter is exclusive with the `path`, `content` `spark_conf` and `libraries` parameters."
             )
-        path = path or make_notebook(content=content)
         name = name or f"dummy-j{make_random(4)}"
         tags = tags or {}
         tags["RemoveAfter"] = tags.get("RemoveAfter", watchdog_remove_after)
@@ -230,8 +236,10 @@ def create(
                 timeout_seconds=0,
             )
             if task_type == SparkPythonTask:
+                path = path or make_workspace_file(content=content)
                 task.spark_python_task = SparkPythonTask(python_file=str(path))
             else:
+                path = path or make_notebook(content=content)
                 task.notebook_task = NotebookTask(notebook_path=str(path))
             tasks = [task]
         response = ws.jobs.create(name=name, tasks=tasks, tags=tags)

tests/unit/fixtures/test_compute.py

@@ -62,16 +62,19 @@ def test_make_job_with_content() -> None:
     tasks = job.settings.tasks
     assert len(tasks) == 1
     workspace_path = WorkspacePath(ctx["ws"], tasks[0].notebook_task.notebook_path)
+    assert not workspace_path.suffix  # Notebooks have no suffix
     assert workspace_path.read_text() == "print(2)"
 
 
 def test_make_job_with_spark_python_task() -> None:
-    _, job = call_stateful(make_job, path="test.py", task_type=SparkPythonTask)
+    ctx, job = call_stateful(make_job, content="print(3)", task_type=SparkPythonTask)
     tasks = job.settings.tasks
     assert len(tasks) == 1
     assert tasks[0].notebook_task is None
     assert tasks[0].spark_python_task is not None
-    assert tasks[0].spark_python_task.python_file == "test.py"
+    workspace_path = WorkspacePath(ctx["ws"], tasks[0].spark_python_task.python_file)
+    assert workspace_path.suffix == ".py"  # Python files have suffix
+    assert workspace_path.read_text() == "print(3)"
 
 
 def test_make_job_with_spark_conf() -> None: