Skip to content

Commit 3aa8265

Browse files
authored
[FEATURE] Extend make_job to run SparkPythonTask (#60)
## Changes Extend `make_job` to run `SparkPythonTask` by: - Add `make_workspace_file` fixture to create the Python file to run. - Support SQL notebooks and files - Add `task_type` to `make_job` to signal what type of task to run (notebook or python file) - Add `instance_pool_id` to `make_job` for speeding up a job in integration tests by reusing an instance pool - Add unit tests for changed fixtures - Run jobs in integration tests ### Tests - [x] manually tested - [x] added unit tests - [x] added integration tests
1 parent 0975eb7 commit 3aa8265

File tree

9 files changed

+404
-82
lines changed

9 files changed

+404
-82
lines changed

README.md

Lines changed: 42 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,7 @@ def test_workspace_operations(ws):
278278
assert len(clusters) >= 0
279279
```
280280

281-
See also [`log_workspace_link`](#log_workspace_link-fixture), [`make_alert_permissions`](#make_alert_permissions-fixture), [`make_authorization_permissions`](#make_authorization_permissions-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_permissions`](#make_cluster_permissions-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_cluster_policy_permissions`](#make_cluster_policy_permissions-fixture), [`make_dashboard_permissions`](#make_dashboard_permissions-fixture), [`make_directory`](#make_directory-fixture), [`make_directory_permissions`](#make_directory_permissions-fixture), [`make_experiment`](#make_experiment-fixture), [`make_experiment_permissions`](#make_experiment_permissions-fixture), [`make_feature_table`](#make_feature_table-fixture), [`make_feature_table_permissions`](#make_feature_table_permissions-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_instance_pool_permissions`](#make_instance_pool_permissions-fixture), [`make_job`](#make_job-fixture), [`make_job_permissions`](#make_job_permissions-fixture), [`make_lakeview_dashboard_permissions`](#make_lakeview_dashboard_permissions-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_notebook_permissions`](#make_notebook_permissions-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_pipeline_permissions`](#make_pipeline_permissions-fixture), [`make_query`](#make_query-fixture), [`make_query_permissions`](#make_query_permissions-fixture), [`make_registered_model_permissions`](#make_registered_model_permissions-fixture), [`make_repo`](#make_repo-fixture), [`make_repo_permissions`](#make_repo_permissions-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`spark`](#spark-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture).
281+
See also [`log_workspace_link`](#log_workspace_link-fixture), [`make_alert_permissions`](#make_alert_permissions-fixture), [`make_authorization_permissions`](#make_authorization_permissions-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_permissions`](#make_cluster_permissions-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_cluster_policy_permissions`](#make_cluster_policy_permissions-fixture), [`make_dashboard_permissions`](#make_dashboard_permissions-fixture), [`make_directory`](#make_directory-fixture), [`make_directory_permissions`](#make_directory_permissions-fixture), [`make_experiment`](#make_experiment-fixture), [`make_experiment_permissions`](#make_experiment_permissions-fixture), [`make_feature_table`](#make_feature_table-fixture), [`make_feature_table_permissions`](#make_feature_table_permissions-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_instance_pool_permissions`](#make_instance_pool_permissions-fixture), [`make_job`](#make_job-fixture), [`make_job_permissions`](#make_job_permissions-fixture), [`make_lakeview_dashboard_permissions`](#make_lakeview_dashboard_permissions-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_notebook_permissions`](#make_notebook_permissions-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_pipeline_permissions`](#make_pipeline_permissions-fixture), [`make_query`](#make_query-fixture), [`make_query_permissions`](#make_query_permissions-fixture), [`make_registered_model_permissions`](#make_registered_model_permissions-fixture), [`make_repo`](#make_repo-fixture), [`make_repo_permissions`](#make_repo_permissions-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file`](#make_workspace_file-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`spark`](#spark-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture).
282282

283283

284284
[[back to top](#python-testing-for-databricks)]
@@ -372,7 +372,7 @@ random_string = make_random(k=8)
372372
assert len(random_string) == 8
373373
```
374374

375-
See also [`make_acc_group`](#make_acc_group-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_feature_table`](#make_feature_table-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture).
375+
See also [`make_acc_group`](#make_acc_group-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_feature_table`](#make_feature_table-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_workspace_file`](#make_workspace_file-fixture).
376376

377377

378378
[[back to top](#python-testing-for-databricks)]
@@ -410,22 +410,26 @@ Create a Databricks job and clean it up after the test. Returns a function to cr
410410
a [`Job`](https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/jobs.html#databricks.sdk.service.jobs.Job) instance.
411411

412412
Keyword Arguments:
413-
* `notebook_path` (str, optional): The path to the notebook. If not provided, a random notebook will be created.
414413
* `name` (str, optional): The name of the job. If not provided, a random name will be generated.
415-
* `spark_conf` (dict, optional): The Spark configuration of the job.
414+
* `path` (str, optional): The path to the notebook or file used in the job. If not provided, a random notebook or file will be created
415+
* [DEPRECATED: Use `path` instead] `notebook_path` (str, optional): The path to the notebook. If not provided, a random notebook will be created.
416+
* `content` (str | bytes, optional): The content of the notebook or file used in the job. If not provided, default content of `make_notebook` will be used.
417+
* `task_type` (type[NotebookTask] | type[SparkPythonTask], optional): The type of task. If not provides, `type[NotebookTask]` will be used.
418+
* `instance_pool_id` (str, optional): The instance pool id to add to the job cluster. If not provided, no instance pool will be used.
419+
* `spark_conf` (dict, optional): The Spark configuration of the job. If not provided, Spark configuration is not explicitly set.
416420
* `libraries` (list, optional): The list of libraries to install on the job.
417-
* other arguments are passed to `WorkspaceClient.jobs.create` method.
418-
419-
If no task argument is provided, a single task with a notebook task will be created, along with a disposable notebook.
420-
Latest Spark version and a single worker clusters will be used to run this ephemeral job.
421+
* `tags` (list[str], optional): A list of job tags. If not provided, no additional tags will be set on the job.
422+
* `tasks` (list[Task], optional): A list of job tags. If not provided, a single task with a notebook task will be
423+
created, along with a disposable notebook. Latest Spark version and a single worker clusters will be used to run
424+
this ephemeral job.
421425

422426
Usage:
423427
```python
424428
def test_job(make_job):
425429
logger.info(f"created {make_job()}")
426430
```
427431

428-
See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`make_notebook`](#make_notebook-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
432+
See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`make_notebook`](#make_notebook-fixture), [`make_workspace_file`](#make_workspace_file-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
429433

430434

431435
[[back to top](#python-testing-for-databricks)]
@@ -621,8 +625,9 @@ The function returns [`os.PathLike` object](https://github.com/databrickslabs/bl
621625

622626
Keyword arguments:
623627
* `path` (str, optional): The path of the notebook. Defaults to `dummy-*` notebook in current user's home folder.
624-
* `content` (typing.BinaryIO, optional): The content of the notebook. Defaults to `print(1)`.
628+
* `content` (str | bytes | io.BinaryIO, optional): The content of the notebook. Defaults to `print(1)` for Python and `SELECT 1` for SQL.
625629
* `language` ([`Language`](https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/workspace.html#databricks.sdk.service.workspace.Language), optional): The language of the notebook. Defaults to `Language.PYTHON`.
630+
* `encoding` (`str`, optional): The file encoding. Defaults to `sys.getdefaultencoding()`.
626631
* `format` ([`ImportFormat`](https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/workspace.html#databricks.sdk.service.workspace.ImportFormat), optional): The format of the notebook. Defaults to `ImportFormat.SOURCE`.
627632
* `overwrite` (bool, optional): Whether to overwrite the notebook if it already exists. Defaults to `False`.
628633

@@ -644,6 +649,32 @@ _No description yet._
644649
See also [`ws`](#ws-fixture).
645650

646651

652+
[[back to top](#python-testing-for-databricks)]
653+
654+
### `make_workspace_file` fixture
655+
Returns a function to create Databricks workspace file and clean up after the test.
656+
The function returns [`os.PathLike` object](https://github.com/databrickslabs/blueprint?tab=readme-ov-file#python-native-pathlibpath-like-interfaces).
657+
658+
Keyword arguments:
659+
* `path` (str, optional): The path of the file. Defaults to `dummy-*` notebook in current user's home folder.
660+
* `content` (str | bytes, optional): The content of the file. Defaults to `print(1)` for Python and `SELECT 1` for SQL.
661+
* `language` ([`Language`](https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/workspace.html#databricks.sdk.service.workspace.Language), optional): The language of the notebook. Defaults to `Language.PYTHON`.
662+
* `encoding` (`str`, optional): The file encoding. Defaults to `sys.getdefaultencoding()`.
663+
664+
This example creates a notebook and verifies that the workspace path is an existing file with contents `print(1)`:
665+
```python
666+
def test_create_file(make_workspace_file):
667+
workspace_file = make_workspace_file()
668+
assert workspace_file.is_file()
669+
assert "print(1)" in workspace_file.read_text()
670+
```
671+
672+
TODO:
673+
Merge functionality with `make_notebook` if `WorkspacePath` supports creating notebooks.
674+
675+
See also [`make_job`](#make_job-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_purge_suffix`](#watchdog_purge_suffix-fixture).
676+
677+
647678
[[back to top](#python-testing-for-databricks)]
648679

649680
### `make_directory` fixture
@@ -1114,7 +1145,7 @@ See also [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster
11141145
### `watchdog_purge_suffix` fixture
11151146
HEX-encoded purge time suffix for test objects.
11161147

1117-
See also [`make_acc_group`](#make_acc_group-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_repo`](#make_repo-fixture), [`make_user`](#make_user-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
1148+
See also [`make_acc_group`](#make_acc_group-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_repo`](#make_repo-fixture), [`make_user`](#make_user-fixture), [`make_workspace_file`](#make_workspace_file-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
11181149

11191150

11201151
[[back to top](#python-testing-for-databricks)]

0 commit comments

Comments
 (0)