pytest-databricks-lakeflow

Test the Lakeflow Jobs in your Databricks Asset Bundle with plain pytest: deploy a throwaway copy of a job or a single task, run it, assert on the result, tear it down.

Why this exists

• A job in a bundle is usually checked by deploying it and pressing Run. This turns that into a pytest test, with assert, parametrize and xdist.
• Every test runs its own copy of the job, so tests run in parallel and never touch the jobs your bundle deployed.
• Task values cannot be read through the Jobs API. The plugin injects tasks that set and capture them, so one task can be tested on its own with the upstream values you choose.

Quickstart

You need Python 3.11 or later, the Databricks CLI v1.7.0 or later on PATH, credentials that both the CLI and the SDK can use (a ~/.databrickscfg profile or the DATABRICKS_* variables), and a workspace with serverless jobs.

uv add --dev pytest-databricks-lakeflow

Or pip install pytest-databricks-lakeflow. Add a test target to databricks.yml. Development mode pauses schedules, and the separate root_path keeps test deployments away from everyone else's:

targets:
  test:
    mode: development
    workspace:
      root_path: ~/.bundle/${bundle.name}/test

Write a test. This one runs my_example_job from examples/bundle; swap in your own job and task keys:

from pytest_databricks_lakeflow import Lakeflow


def test_job_succeeds(lakeflow: Lakeflow) -> None:
    with lakeflow.job("my_example_job", params={"param1": "hello", "param2": "world"}) as job:
        run = job.run()
        assert run.succeeded
        assert run.task("transform").output == "hello world"
        assert "Success! hello" in run.task("my_wheel_task").logs

Run pytest. The first test deploys the bundle's test target once, then every test creates, runs and deletes its own job. A failing test keeps its job, and the job and run URLs appear in the test report. Write assert run.succeeded, run and a failure reads:

AssertionError: JobRun('my_example_job', FAILED, https://<workspace>/#job/123/run/456, failing=['transform: PermissionError: forbidden may not read /Volumes/example/raw/landing'])

To change the target, set a catalog for temporary schemas or lengthen the timeout, override the lakeflow fixture in conftest.py:

from datetime import timedelta

import pytest

from pytest_databricks_lakeflow import Lakeflow


@pytest.fixture(scope="session")
def lakeflow() -> Lakeflow:
    return Lakeflow(target="test", catalog="dev_sandbox", timeout=timedelta(minutes=20))

Testing one task

lakeflow.task("job_key.task_key") builds a job that holds only that task. Each upstream task you name in upstream_task_values is replaced by a stub task that sets those task values, and capture_task_values adds a probe task that reports the values the task set, even when it failed.

import pytest

from pytest_databricks_lakeflow import Lakeflow


@pytest.mark.parametrize(("greeting", "expected"), [("hello", "hello world"), ("goodbye", "goodbye world")])
def test_transform(lakeflow: Lakeflow, greeting: str, expected: str) -> None:
    with lakeflow.task(
        "my_example_job.transform",
        params={"param1": greeting, "param2": "world"},
        upstream_task_values={"extract": {"row_count": 42, "source_path": "/Volumes/x/y/z"}},
        capture_task_values=["rows_written"],
    ) as task:
        run = task.run()
        assert run.succeeded
        assert run.task("transform").output == expected
        assert run.task("transform").task_values["rows_written"] == 42

Inside the task, dbutils.jobs.taskValues.get(taskKey="extract", key="source_path") and {{tasks.extract.values.row_count}} resolve against the stub. Reading a value the task never set raises a KeyError that says so.

How it works

1. The first test that needs a job runs databricks bundle deploy for the target once per session and reads each deployed job's settings through the Jobs API.
2. lakeflow.job() copies those settings into a new job named [pytest] <test id>, with schedules, triggers, notifications, run_as and the bundle marker removed.
3. run() starts the job, waits, and returns a JobRun with every task's state, output, logs, error and captured values.
4. Leaving the with block deletes the job unless the test failed. Session end cancels anything still running.
5. lakeflow.task() does the same for one task, with the stub and probe tasks on serverless compute or a job cluster you name.

Documentation

Everything else is on the docs site: https://stevenayers.github.io/pytest-databricks-lakeflow/

• Getting started: prerequisites, install, bundle target, first test
• Writing job tests: parameters, only, repeated runs, failures, temporary schemas, shared runs
• Writing task tests: stubs, probes, value rules, helper compute
• Configuration: every Lakeflow argument, command-line flag and sync strategy
• How it works: the deploy and teardown lifecycle, xdist, timing
• Limitations and FAQ: no logs for notebook tasks, what probes cost, what stays in the workspace
• API reference

Contributing and licence

Bug reports and pull requests are welcome. See CONTRIBUTING.md for the development setup, the test suites and how releases work, and SUPPORT.md for where to ask questions.

Apache-2.0. See LICENSE.