pytest-databricks-lakeflow 0.1.0__py3-none-any.whl

pytest_databricks_lakeflow/__init__.py

@@ -0,0 +1,20 @@
+"""Run Lakeflow Jobs from a Databricks Asset Bundle as throwaway jobs inside pytest tests."""
+
+from __future__ import annotations
+
+from pytest_databricks_lakeflow.client import Lakeflow
+from pytest_databricks_lakeflow.ephemeral import EphemeralJob
+from pytest_databricks_lakeflow.errors import LakeflowBundleError, LakeflowError, LakeflowTimeout
+from pytest_databricks_lakeflow.results import JobRun, TaskError, TaskRun, TaskValues
+
+__all__ = [
+    "EphemeralJob",
+    "JobRun",
+    "Lakeflow",
+    "LakeflowBundleError",
+    "LakeflowError",
+    "LakeflowTimeout",
+    "TaskError",
+    "TaskRun",
+    "TaskValues",
+]

pytest_databricks_lakeflow/bundle.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import copy
+import glob
+import json
+import logging
+import re
+import subprocess
+from collections.abc import Callable, Mapping
+from pathlib import Path
+from typing import Any
+
+from databricks.sdk.service import jobs
+
+from pytest_databricks_lakeflow.errors import LakeflowBundleError
+
+log = logging.getLogger("pytest_databricks_lakeflow")
+
+BUNDLE_FILES = ("databricks.yml", "databricks.yaml")
+MIN_CLI_VERSION = (1, 7, 0)
+UNRESOLVED = re.compile(r"\$\{[^}]+\}")
+VERSION = re.compile(r"v(\d+)\.(\d+)\.(\d+)")
+CLI_DOCS = "https://docs.databricks.com/dev-tools/cli/install"
+
+
+def find_root(start: Path) -> Path:
+    for path in (start, *start.parents):
+        if _has_bundle(path):
+            return path
+    raise LakeflowBundleError(f"No databricks.yml in {start} or any parent directory. Set lakeflow_bundle_root.")
+
+
+def require_root(path: Path) -> Path:
+    if not _has_bundle(path):
+        raise LakeflowBundleError(f"lakeflow_bundle_root {path} has no databricks.yml")
+    return path
+
+
+def _has_bundle(path: Path) -> bool:
+    return any((path / name).is_file() for name in BUNDLE_FILES)
+
+
+def run_cli(args: list[str], root: Path, profile: str | None = None) -> str:
+    if not root.is_dir():
+        raise LakeflowBundleError(f"Bundle root {root} does not exist")
+    cmd = ["databricks", *args, *(["--profile", profile] if profile else [])]
+    try:
+        result = subprocess.run(cmd, cwd=root, check=True, capture_output=True, text=True)
+    except FileNotFoundError:
+        raise LakeflowBundleError(f"The databricks CLI is not on PATH. Install it from {CLI_DOCS}") from None
+    except subprocess.CalledProcessError as e:
+        raise LakeflowBundleError(f"`{' '.join(cmd)}` failed with exit code {e.returncode}:\n{e.stderr.strip()}") from e
+    return result.stdout
+
+
+def check_cli_version(root: Path) -> None:
+    out = run_cli(["--version"], root)
+    match = VERSION.search(out)
+    version = tuple(int(part) for part in match.groups()) if match else None
+    if version and (0, 0, 0) < version < MIN_CLI_VERSION:
+        minimum = ".".join(map(str, MIN_CLI_VERSION))
+        raise LakeflowBundleError(f"pytest-databricks-lakeflow needs Databricks CLI v{minimum} or later: {out.strip()}")
+
+
+def deploy(root: Path, target: str, profile: str | None) -> None:
+    log.info("Deploying bundle %s to target %s", root, target)
+    run_cli(["bundle", "deploy", "-t", target, "--auto-approve"], root, profile)
+
+
+def sync(root: Path, target: str, profile: str | None) -> None:
+    log.info("Syncing bundle %s to target %s", root, target)
+    run_cli(["bundle", "sync", "-t", target], root, profile)
+
+
+def validate(root: Path, target: str, profile: str | None) -> dict[str, Any]:
+    return _json_command("validate", root, target, profile)
+
+
+def summary(root: Path, target: str, profile: str | None) -> dict[str, Any]:
+    return _json_command("summary", root, target, profile)
+
+
+def destroy(root: Path, target: str, profile: str | None) -> None:
+    log.info("Destroying bundle %s target %s", root, target)
+    run_cli(["bundle", "destroy", "-t", target, "--auto-approve"], root, profile)
+
+
+def _json_command(command: str, root: Path, target: str, profile: str | None) -> dict[str, Any]:
+    out = run_cli(["bundle", command, "-t", target, "-o", "json"], root, profile)
+    try:
+        data = json.loads(out)
+    except json.JSONDecodeError as e:
+        raise LakeflowBundleError(f"`databricks bundle {command}` did not print JSON: {e}\n{out[:2000]}") from e
+    if not isinstance(data, dict):
+        raise LakeflowBundleError(f"`databricks bundle {command}` printed {type(data).__name__}, expected an object")
+    return data
+
+
+def job_ids(bundle_summary: Mapping[str, Any]) -> tuple[dict[str, int], dict[str, str]]:
+    """Read deployed job ids from `bundle summary` output.
+
+    Args:
+        bundle_summary: The parsed output of `databricks bundle summary -o json`.
+
+    Returns:
+        The job ids by resource key, and an error for each job that is defined but not deployed.
+    """
+    ids: dict[str, int] = {}
+    errors: dict[str, str] = {}
+    for key, job in bundle_summary.get("resources", {}).get("jobs", {}).items():
+        if job.get("modified_status") == "deleted":
+            continue
+        if job.get("id"):
+            ids[key] = int(job["id"])
+        else:
+            errors[key] = f"Job {key!r} is not deployed. Run `databricks bundle deploy` or use lakeflow_sync = deploy."
+    return ids, errors
+
+
+def templates_from_validate(
+    config: Mapping[str, Any], root: Path, strategy: str, exists: Callable[[str], bool]
+) -> tuple[dict[str, dict[str, Any]], dict[str, str]]:
+    """Build job templates from `bundle validate` output, for the sync and none strategies.
+
+    Validate resolves notebook and script paths but leaves library paths local. With `none`, each local library
+    is pointed at the copy that an earlier `bundle deploy` uploaded, after checking that it is there. With `sync`,
+    which uploads no libraries, a job with local libraries gets an error instead.
+
+    Args:
+        config: The parsed output of `databricks bundle validate -o json`.
+        root: The bundle root, against which local library paths are resolved.
+        strategy: `sync` or `none`.
+        exists: Checks whether a workspace or Volumes path exists.
+
+    Returns:
+        The templates by resource key, and an error for each job that cannot be used.
+    """
+    artifact_path = config.get("workspace", {}).get("artifact_path", "")
+    templates: dict[str, dict[str, Any]] = {}
+    errors: dict[str, str] = {}
+    for key, job in config.get("resources", {}).get("jobs", {}).items():
+        refs = sorted(set(UNRESOLVED.findall(json.dumps(job))))
+        if refs:
+            errors[key] = (
+                f"Job {key!r} contains {', '.join(refs)}, which `bundle validate` does not resolve. "
+                "Use lakeflow_sync = deploy."
+            )
+            continue
+        local = _local_libraries(job)
+        if local and strategy == "sync":
+            errors[key] = (
+                f"Job {key!r} uses local libraries ({', '.join(local)}), which `bundle sync` does not upload. "
+                "Use lakeflow_sync = deploy."
+            )
+            continue
+        patched = job
+        if local:
+            try:
+                patched = _map_libraries(job, lambda path: _uploaded(path, root, artifact_path, exists))
+            except LakeflowBundleError as e:
+                errors[key] = f"Job {key!r}: {e}"
+                continue
+        templates[key] = jobs.JobSettings.from_dict(patched).as_dict()
+    return templates, errors
+
+
+def _uploaded(pattern: str, root: Path, artifact_path: str, exists: Callable[[str], bool]) -> list[str]:
+    matches = sorted(glob.glob(str(root / pattern)))
+    if not matches:
+        raise LakeflowBundleError(f"no file under {root} matches {pattern!r}. Build and deploy the bundle first.")
+    remote = [f"{artifact_path}/.internal/{Path(match).name}" for match in matches]
+    missing = [path for path in remote if not exists(path)]
+    if missing:
+        raise LakeflowBundleError(f"{', '.join(missing)} not found. Run `databricks bundle deploy` first.")
+    return remote
+
+
+def _local_libraries(job: Mapping[str, Any]) -> list[str]:
+    found: list[str] = []
+
+    def record(path: str) -> list[str]:
+        found.append(path)
+        return [path]
+
+    _map_libraries(job, record)
+    return found
+
+
+def _map_libraries(job: Mapping[str, Any], fn: Callable[[str], list[str]]) -> dict[str, Any]:
+    """Copy a job, replacing each local library path with the paths `fn` returns for it."""
+    copied: dict[str, Any] = copy.deepcopy(dict(job))
+    tasks = [*copied.get("tasks", [])]
+    tasks += [t["for_each_task"]["task"] for t in tasks if "task" in t.get("for_each_task", {})]
+    for task in tasks:
+        if "libraries" not in task:
+            continue
+        libraries = []
+        for library in task["libraries"]:
+            kind = next((k for k in ("whl", "jar") if _is_local(library.get(k))), None)
+            libraries += [{kind: path} for path in fn(library[kind])] if kind else [library]
+        task["libraries"] = libraries
+    for env in copied.get("environments", []):
+        spec = env.get("spec", {})
+        if "dependencies" in spec:
+            spec["dependencies"] = [path for dep in spec["dependencies"] for path in _dependency(dep, fn)]
+    return copied
+
+
+def _dependency(dep: str, fn: Callable[[str], list[str]]) -> list[str]:
+    local = dep.startswith(("./", "../")) or (dep.endswith((".whl", ".jar")) and _is_local(dep))
+    return fn(dep) if local else [dep]
+
+
+def _is_local(path: object) -> bool:
+    return isinstance(path, str) and bool(path) and not path.startswith("/") and ":" not in path

pytest_databricks_lakeflow/client.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+import copy
+import logging
+import time
+import warnings
+from collections.abc import Collection, Mapping, Sequence
+from dataclasses import dataclass, field, replace
+from typing import Any
+
+import pytest
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors.base import DatabricksError
+from databricks.sdk.service import jobs
+
+from pytest_databricks_lakeflow.ephemeral import EphemeralJob, cancel_run, item_failed
+from pytest_databricks_lakeflow.errors import LakeflowBundleError, LakeflowError
+from pytest_databricks_lakeflow.options import Options
+from pytest_databricks_lakeflow.settings import Helpers, check_params, inject_probes, sanitise, wrap_task
+
+log = logging.getLogger("pytest_databricks_lakeflow")
+
+
+@dataclass
+class Lakeflow:
+    """Creates throwaway copies of the bundle's jobs. Use the `lakeflow` fixture to get one.
+
+    Attributes:
+        ws: The workspace client.
+        options: The plugin's options.
+        run_id: The id of this test session, used to tag ephemeral jobs.
+        templates: The deployed settings of each bundle job, by resource key, as dicts.
+        template_errors: Why each bundle job that cannot be used has no template, by resource key.
+        helpers: Workspace paths of the helper notebooks, by name.
+        request: The test this copy is bound to, or `None` for the session's copy.
+        active_runs: Ids of runs started in this session that have not finished.
+    """
+
+    ws: WorkspaceClient
+    options: Options
+    run_id: str
+    templates: dict[str, dict[str, Any]]
+    template_errors: dict[str, str] = field(default_factory=dict)
+    helpers: dict[str, str] = field(default_factory=dict)
+    request: pytest.FixtureRequest | None = None
+    active_runs: set[int] = field(default_factory=set)
+
+    @classmethod
+    def from_state(cls, ws: WorkspaceClient, options: Options, state: dict[str, Any]) -> Lakeflow:
+        return cls(
+            ws=ws,
+            options=options,
+            run_id=state["run_id"],
+            templates=state["templates"],
+            template_errors=state["template_errors"],
+            helpers=state["helpers"],
+        )
+
+    def for_test(self, request: pytest.FixtureRequest) -> Lakeflow:
+        """Bind a copy to one test.
+
+        Args:
+            request: The test's fixture request.
+
+        Returns:
+            A copy that shares this one's templates and active runs.
+        """
+        return replace(self, request=request)
+
+    def template(self, key: str) -> jobs.JobSettings:
+        """Get a fresh copy of the deployed settings of a bundle job.
+
+        Args:
+            key: The bundle resource key of the job.
+
+        Returns:
+            The job's settings, safe to change.
+
+        Raises:
+            KeyError: If the bundle has no job `key`.
+            LakeflowBundleError: If the job cannot be used as a template, for example because it is not deployed.
+        """
+        if key in self.template_errors:
+            raise LakeflowBundleError(self.template_errors[key])
+        if key not in self.templates:
+            raise KeyError(f"No job {key!r} in the bundle. Available: {sorted(self.templates)}")
+        return jobs.JobSettings.from_dict(copy.deepcopy(self.templates[key]))
+
+    def current_request(self) -> pytest.FixtureRequest:
+        if self.request is None:
+            raise LakeflowError("Use the lakeflow fixture inside a test to create ephemeral jobs")
+        return self.request
+
+    def job(
+        self,
+        key: str,
+        params: Mapping[str, str] | None = None,
+        capture_values: Mapping[str, Sequence[str]] | None = None,
+        name: str | None = None,
+    ) -> EphemeralJob:
+        """Make a throwaway copy of a bundle job.
+
+        Args:
+            key: The bundle resource key of the job, as `databricks bundle run` takes it.
+            params: Job parameters for every run of the copy.
+            capture_values: The task value keys to capture, by task key.
+            name: Replaces the test node id in the job name.
+
+        Returns:
+            The job, created when entered with `with` or on its first run.
+
+        Raises:
+            KeyError: If the bundle has no job `key`.
+            LakeflowError: If `capture_values` names a task that is not in the job, or a parameter name is
+                reserved.
+        """
+        settings = self.template(key)
+        check_params(key, settings, params or {}, allow_undeclared=True)
+        inject_probes(settings, capture_values or {}, self._helpers())
+        return self._ephemeral(key, settings, params, capture_values, (), name)
+
+    def task(
+        self,
+        ref: str,
+        params: Mapping[str, str] | None = None,
+        upstream_values: Mapping[str, Mapping[str, object]] | None = None,
+        capture_values: Sequence[str] | None = None,
+    ) -> EphemeralJob:
+        """Make a throwaway job that holds one task of a bundle job.
+
+        Each upstream task named in `upstream_values` is replaced by a stub task with the same key, so the task's
+        `{{tasks.<key>.values.<name>}}` references and `taskValues.get` calls resolve against the stub.
+
+        Args:
+            ref: The job and task as `"job_key.task_key"`.
+            params: Job parameters for every run.
+            upstream_values: The task values each stub sets, by upstream task key.
+            capture_values: The task value keys to capture from the task.
+
+        Returns:
+            The job, created when entered with `with` or on its first run.
+
+        Raises:
+            KeyError: If the bundle has no such job.
+            LakeflowError: If the job has no such task, the task refers to task values of a task that is not
+                stubbed, or a value is too large or not JSON-serialisable.
+        """
+        job_key, _, task_key = ref.partition(".")
+        if not job_key or not task_key:
+            raise LakeflowError(f"Expected 'job_key.task_key', got {ref!r}")
+        settings = self.template(job_key)
+        check_params(job_key, settings, params or {}, allow_undeclared=True)
+        upstream = dict(upstream_values or {})
+        wrap_task(settings, task_key, upstream, list(capture_values or []), self._helpers())
+        capture = {task_key: list(capture_values)} if capture_values else {}
+        return self._ephemeral(job_key, settings, params, capture, upstream, None)
+
+    def _helpers(self) -> Helpers:
+        return Helpers(
+            set_values=self.helpers.get("set_values", ""),
+            get_values=self.helpers.get("get_values", ""),
+            environment_version=self.options.helpers_environment_version,
+            job_cluster_key=self.options.helpers_compute,
+        )
+
+    def _ephemeral(
+        self,
+        key: str,
+        settings: jobs.JobSettings,
+        params: Mapping[str, str] | None,
+        capture_values: Mapping[str, Sequence[str]] | None,
+        stubs: Collection[str],
+        name: str | None,
+    ) -> EphemeralJob:
+        request = self.current_request()
+        if settings.performance_target == jobs.PerformanceTarget.STANDARD and not self.options.performance_optimized:
+            warnings.warn(
+                f"Job {key!r} has performance_target STANDARD, which adds 3 to 4 minutes to every run. "
+                "Set lakeflow_performance_optimized = true to run tests with PERFORMANCE_OPTIMIZED.",
+                stacklevel=3,
+            )
+        nodeid = request.node.nodeid
+        sanitise(
+            settings,
+            name=name or nodeid,
+            tags={"pytest": "", "pytest_run_id": self.run_id, "pytest_nodeid": nodeid},
+            disable_auto_optimization=self.options.disable_auto_optimization,
+            performance_optimized=self.options.performance_optimized,
+        )
+        job = EphemeralJob(self, key, settings, params, capture_values, stubs)
+        request.addfinalizer(lambda: job.close(failed=item_failed(request.node.stash)))
+        return job
+
+    def sweep(self) -> list[int]:
+        """Delete this user's ephemeral jobs from other runs that are older than `lakeflow_sweep_hours`.
+
+        Returns:
+            The ids of the deleted jobs.
+        """
+        me = self.ws.current_user.me().user_name
+        cutoff = (time.time() - self.options.sweep_hours * 3600) * 1000
+        stale = [
+            job.job_id
+            for job in self.ws.jobs.list(limit=100)
+            if job.job_id is not None and job.creator_user_name == me and _stale(job, self.run_id, cutoff)
+        ]
+        deleted = []
+        for job_id in stale:
+            try:
+                self.ws.jobs.delete(job_id)
+            except DatabricksError as e:
+                log.warning("Could not delete job %s: %s", job_id, e)
+            else:
+                deleted.append(job_id)
+        log.info("Swept %d leftover ephemeral jobs", len(deleted))
+        return deleted
+
+    def cancel_active_runs(self) -> None:
+        """Cancel every run started in this session that has not finished."""
+        for run_id in sorted(self.active_runs):
+            cancel_run(self.ws, run_id)
+        self.active_runs.clear()
+
+
+def _stale(job: jobs.BaseJob, run_id: str, cutoff_ms: float) -> bool:
+    tags = (job.settings.tags if job.settings else None) or {}
+    other_run = tags.get("pytest_run_id", run_id) != run_id
+    return "pytest" in tags and other_run and (job.created_time or 0) < cutoff_ms

pytest_databricks_lakeflow/ephemeral.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Collection, Mapping, Sequence
+from datetime import timedelta
+from types import TracebackType
+from typing import TYPE_CHECKING, cast
+
+import pytest
+from databricks.sdk.errors.base import DatabricksError
+from databricks.sdk.service import jobs
+
+from pytest_databricks_lakeflow.errors import LakeflowError, LakeflowTimeout
+from pytest_databricks_lakeflow.results import JobRun, build_job_run
+from pytest_databricks_lakeflow.settings import check_params, create_kwargs, select
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+    from pytest_databricks_lakeflow.client import Lakeflow
+
+log = logging.getLogger("pytest_databricks_lakeflow")
+
+POLL_SECONDS = 2.0
+REPORTS = pytest.StashKey[dict[str, pytest.TestReport]]()
+
+
+class EphemeralJob:
+    """A throwaway copy of a bundle job.
+
+    The job is created on entering the `with` block, or by the first `run()` when used without `with`. It is
+    deleted on leaving the block, or at the end of the test, unless `--lakeflow-keep` is set or the test failed
+    and `--lakeflow-keep-failed` is on.
+
+    Attributes:
+        lakeflow: The `Lakeflow` that made the job.
+        job_key: The bundle resource key of the job it copies.
+        settings: The settings the job is created with.
+        params: Job parameters for every run.
+        capture_values: The task value keys captured by probes, by task key.
+        stubs: Keys of the stub tasks.
+        job_id: The job's id, once created.
+        runs: The outcome of each finished run, in order.
+    """
+
+    def __init__(
+        self,
+        lakeflow: Lakeflow,
+        job_key: str,
+        settings: jobs.JobSettings,
+        params: Mapping[str, str] | None = None,
+        capture_values: Mapping[str, Sequence[str]] | None = None,
+        stubs: Collection[str] = (),
+    ) -> None:
+        self.lakeflow = lakeflow
+        self.job_key = job_key
+        self.settings = settings
+        self.params = dict(params or {})
+        self.capture_values = {key: list(keys) for key, keys in (capture_values or {}).items()}
+        self.stubs = frozenset(stubs)
+        self.job_id: int | None = None
+        self.runs: list[JobRun] = []
+        self._run_urls: dict[int, str] = {}
+        self._active: set[int] = set()
+        self._closed = False
+
+    @property
+    def url(self) -> str | None:
+        """The job's page in the workspace, or `None` before it is created."""
+        return None if self.job_id is None else f"{self.lakeflow.ws.config.host}/jobs/{self.job_id}"
+
+    def __enter__(self) -> EphemeralJob:
+        self.create()
+        return self
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
+    ) -> None:
+        skipped = exc_type is not None and issubclass(exc_type, (pytest.skip.Exception, pytest.xfail.Exception))
+        self.close(failed=exc_type is not None and not skipped)
+
+    def create(self) -> None:
+        """Create the job if it does not exist yet.
+
+        Raises:
+            LakeflowError: If the job has been closed, or `params` names a parameter that the job does not
+                declare after the `pytest_lakeflow_modify_settings` hook has run.
+        """
+        if self._closed:
+            raise LakeflowError(f"The ephemeral job for {self.job_key!r} has already been closed")
+        if self.job_id is not None:
+            return
+        request = self.lakeflow.current_request()
+        request.node.ihook.pytest_lakeflow_modify_settings(settings=self.settings, request=request)
+        check_params(self.job_key, self.settings, self.params, self.lakeflow.options.allow_undeclared_params)
+        self.job_id = self.lakeflow.ws.jobs.create(**create_kwargs(self.settings)).job_id
+        log.info("Created job %s from %s: %s", self.job_id, self.job_key, self.url)
+        request.node.user_properties.append(("lakeflow_job_url", self.url))
+
+    def run(
+        self,
+        params: Mapping[str, str] | None = None,
+        only: Sequence[str] | None = None,
+        timeout: timedelta | None = None,
+    ) -> JobRun:
+        """Run the job and wait for it to finish, creating it first if needed.
+
+        Args:
+            params: Job parameters for this run only, merged over the job's params.
+            only: Task keys to run. `+key`, `key+` and `+key+` add upstream and downstream tasks.
+            timeout: How long to wait. Defaults to `lakeflow_timeout_minutes`.
+
+        Returns:
+            The run's outcome.
+
+        Raises:
+            LakeflowError: If the job has been closed, a parameter is not declared, or `only` names an unknown
+                task.
+            LakeflowTimeout: If the run did not finish in time. The run is cancelled first.
+        """
+        self.create()
+        assert self.job_id is not None
+        options = self.lakeflow.options
+        check_params(self.job_key, self.settings, params or {}, options.allow_undeclared_params)
+        merged = {**self.params, **(params or {})}
+        selected = select(self.settings, only, self.stubs) if only else None
+        ws = self.lakeflow.ws
+        waiter = ws.jobs.run_now(job_id=self.job_id, job_parameters=merged or None, only=selected)
+        run_id = cast(jobs.RunNowResponse, waiter.response).run_id
+        assert run_id is not None
+        self._active.add(run_id)
+        self.lakeflow.active_runs.add(run_id)
+        log.info("Started run %s of job %s", run_id, self.job_id)
+        run = self._wait(run_id, options.timeout if timeout is None else timeout)
+        self._forget(run_id)
+        result = build_job_run(
+            ws, run, job_key=self.job_key, job_id=self.job_id, stubs=self.stubs, capture_values=self.capture_values
+        )
+        self.runs.append(result)
+        log.info("Run %s finished %s in %s: %s", run_id, result.result_state, result.duration, result.run_url)
+        request = self.lakeflow.current_request()
+        request.node.ihook.pytest_lakeflow_run_finished(run=result, request=request)
+        return result
+
+    def _wait(self, run_id: int, timeout: timedelta) -> jobs.Run:
+        ws = self.lakeflow.ws
+        deadline = time.monotonic() + timeout.total_seconds()
+        while True:
+            run = ws.jobs.get_run(run_id)
+            if run_id not in self._run_urls and run.run_page_url:
+                self._run_urls[run_id] = run.run_page_url
+                self.lakeflow.current_request().node.user_properties.append(("lakeflow_run_url", run.run_page_url))
+            if _terminated(run):
+                return run
+            if time.monotonic() >= deadline:
+                self._cancel(run_id)
+                raise LakeflowTimeout(
+                    f"Run {run_id} of {self.job_key!r} did not finish within {timeout} and was cancelled: "
+                    f"{run.run_page_url}"
+                )
+            time.sleep(POLL_SECONDS)
+
+    def close(self, failed: bool = False) -> None:
+        """Cancel active runs, then delete the job unless the keep rules say to keep it.
+
+        Args:
+            failed: Whether the test failed, which keeps the job when `--lakeflow-keep-failed` is on.
+        """
+        if self._closed:
+            return
+        self._closed = True
+        for run_id in sorted(self._active):
+            self._cancel(run_id)
+        if self.job_id is None:
+            return
+        options = self.lakeflow.options
+        if options.keep or (failed and options.keep_failed):
+            log.warning("Keeping job %s from %s: %s", self.job_id, self.job_key, self.url)
+            for url in self._run_urls.values():
+                log.warning("Kept run: %s", url)
+            return
+        self.lakeflow.ws.jobs.delete(self.job_id)
+        log.info("Deleted job %s", self.job_id)
+
+    def _cancel(self, run_id: int) -> None:
+        cancel_run(self.lakeflow.ws, run_id)
+        self._forget(run_id)
+
+    def _forget(self, run_id: int) -> None:
+        self._active.discard(run_id)
+        self.lakeflow.active_runs.discard(run_id)
+
+
+def cancel_run(ws: WorkspaceClient, run_id: int) -> None:
+    try:
+        ws.jobs.cancel_run(run_id)
+        log.info("Cancelled run %s", run_id)
+    except DatabricksError as e:
+        log.warning("Could not cancel run %s: %s", run_id, e)
+
+
+def _terminated(run: jobs.Run) -> bool:
+    if run.status and run.status.state:
+        return run.status.state == jobs.RunLifecycleStateV2State.TERMINATED
+    state = run.state.life_cycle_state if run.state else None
+    return state in (
+        jobs.RunLifeCycleState.TERMINATED,
+        jobs.RunLifeCycleState.SKIPPED,
+        jobs.RunLifeCycleState.INTERNAL_ERROR,
+    )
+
+
+def item_failed(stash: pytest.Stash) -> bool:
+    return any(report.failed for report in stash.get(REPORTS, {}).values())

pytest_databricks_lakeflow/errors.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+
+class LakeflowError(Exception):
+    """Base class for errors raised by pytest-databricks-lakeflow."""
+
+
+class LakeflowBundleError(LakeflowError):
+    """A Databricks CLI bundle command failed, or the bundle cannot supply a job template."""
+
+
+class LakeflowTimeout(LakeflowError):
+    """A job run did not finish within its timeout. The run has been cancelled."""