Python Testing for Databricks

<!-- TOC -->

• Python Testing for Databricks
  • Installation
  • Ecosystem
  • PyTest Fixtures
    • Logging
    • debug_env_name fixture
    • debug_env fixture
    • env_or_skip fixture
    • ws fixture
    • make_run_as fixture
    • acc fixture
    • spark fixture
    • sql_backend fixture
    • sql_exec fixture
    • sql_fetch_all fixture
    • make_random fixture
    • make_instance_pool fixture
    • make_instance_pool_permissions fixture
    • make_job fixture
    • make_job_permissions fixture
    • make_cluster fixture
    • make_cluster_permissions fixture
    • make_cluster_policy fixture
    • make_cluster_policy_permissions fixture
    • make_pipeline fixture
    • make_warehouse fixture
    • make_group fixture
    • make_acc_group fixture
    • make_user fixture
    • make_pipeline_permissions fixture
    • make_notebook fixture
    • make_notebook_permissions fixture
    • make_workspace_file fixture
    • make_directory fixture
    • make_directory_permissions fixture
    • make_repo fixture
    • make_repo_permissions fixture
    • make_workspace_file_permissions fixture
    • make_workspace_file_path_permissions fixture
    • make_secret_scope fixture
    • make_secret_scope_acl fixture
    • make_authorization_permissions fixture
    • make_udf fixture
    • make_catalog fixture
    • make_schema fixture
    • make_table fixture
    • make_storage_credential fixture
    • make_volume fixture
    • product_info fixture
    • make_model fixture
    • make_experiment fixture
    • make_experiment_permissions fixture
    • make_warehouse_permissions fixture
    • make_lakeview_dashboard_permissions fixture
    • log_workspace_link fixture
    • log_account_link fixture
    • make_dashboard_permissions fixture
    • make_alert_permissions fixture
    • make_query fixture
    • make_query_permissions fixture
    • make_registered_model_permissions fixture
    • make_serving_endpoint fixture
    • make_serving_endpoint_permissions fixture
    • make_feature_table fixture
    • make_feature_table_permissions fixture
    • watchdog_remove_after fixture
    • watchdog_purge_suffix fixture
    • is_in_debug fixture
• Project Support

<!-- TOC -->

Installation

Add a databricks-labs-pytester dependency to your pyproject.toml file (or legacy requirements.txt file). You can also install it directly from the command line:

pip install databricks-labs-pytester

If you use hatch as a build system, make sure to add databricks-labs-pytester as a test-time dependency and not as a compile-time dependency, otherwise your wheels will transitively depend on pytest, which is not usually something you need.

[project]
name = "name-of-your-project"
# ...
dependencies = [
  "databricks-sdk~=0.30",
  # ... dependencies required for your code to execute
]

[tool.hatch.envs.default]
dependencies = [
  # ... dependencies required to test/validate/format your code:
    "black~=24.3.0",
    "coverage[toml]~=7.4.4",
    "mypy~=1.9.0",
    "pylint~=3.2.2",
    "pylint-pytest==2.0.0a0",
    "databricks-labs-pylint~=0.4.0",
    "databricks-labs-pytester~=0.2", # <= this library
    "pytest~=8.3.3",
    "pytest-cov~=4.1.0",
    "pytest-mock~=3.14.0",
    "pytest-timeout~=2.3.1",
    "pytest-xdist~=3.5.0",
    "python-lsp-server>=1.9.0",
    "ruff~=0.3.4",
    "types-PyYAML~=6.0.12",
    "types-requests~=2.31.0",
]

[back to top]

Ecosystem

Built on top of Databricks SDK for Python, this library is part of the Databricks Labs Python ecosystem, which includes the following projects:

• PyLint Plugin for Databricks for static code analysis and early bug detection.
• Blueprint for Python-native pathlib.Path-like interfaces, Managing Python App installations within Databricks Workspaces, Application Migrations, and Building Wheels.
• LSQL for lightweight SQL handling and dashboards-as-code.
• UCX for automated migrations into Unity Catalog and LSP plugin for static code analysis for UC compatibility.

See this video for a quick overview of the Databricks Labs Python ecosystem.

[back to top]

PyTest Fixtures

PyTest Fixtures are a powerful way to manage test setup and teardown in Python. This library provides a set of fixtures to help you write integration tests for Databricks. These fixtures were incubated within the Unity Catalog Automated Migrations project for more than a year and are now available for other projects to simplify integration testing with Databricks.

[back to top]

Logging

This library is built on years of debugging integration tests for Databricks and its ecosystem.

That's why it comes with a built-in logger that traces creation and deletion of dummy entities through links in the Databricks Workspace UI. If you run the following code:

def test_new_user(make_user, ws):
    new_user = make_user()
    home_dir = ws.workspace.get_status(f"/Users/{new_user.user_name}")
    assert home_dir.object_type == ObjectType.DIRECTORY

You will see the following output, where the first line is clickable and will take you to the user's profile in the Databricks Workspace UI:

12:30:53  INFO [d.l.p.fixtures.baseline] Created dummy-xwuq-...@example.com: https://.....azuredatabricks.net/#settings/workspace/identity-and-access/users/735...
12:30:53 DEBUG [d.l.p.fixtures.baseline] added workspace user fixture: User(active=True, display_name='dummy-xwuq-...@example.com', ...)
12:30:58 DEBUG [d.l.p.fixtures.baseline] clearing 1 workspace user fixtures
12:30:58 DEBUG [d.l.p.fixtures.baseline] removing workspace user fixture: User(active=True, display_name='dummy-xwuq-...@example.com', ...)

You may need to add the following to your conftest.py file to enable this:

import logging

from databricks.labs.blueprint.logger import install_logger

install_logger()

logging.getLogger('databricks.labs.pytester').setLevel(logging.DEBUG)

[back to top]

<!-- FIXTURES -->

debug_env_name fixture

Specify the name of the debug environment. By default, it is set to .env, which will try to find a file named .env in any of the parent directories of the current working directory and load the environment variables from it via the debug_env fixture.

Alternatively, if you are concerned of the risk of .env files getting checked into version control, we recommend using the ~/.databricks/debug-env.json file to store different sets o