Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,19 @@ def set_utc_timezone():
os.environ.pop("TZ")


@pytest.fixture
def skip_if_runtime_not_serverless_compatible(ws, debug_env):
"""
Skip the test if the cluster runtime does not support serverless clusters.

Args:
ws (WorkspaceClient): Workspace client to interact with Databricks.
debug_env (dict): Test environment variables.
"""
if not os.getenv("DATABRICKS_SERVERLESS_COMPUTE_ID"):
pytest.skip("Skipping test: not running on a serverless cluster")


@pytest.fixture
def skip_if_runtime_not_geo_compatible(ws, debug_env):
"""
Expand Down
59 changes: 31 additions & 28 deletions tests/e2e/test_run_demos.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
from pathlib import Path
from uuid import uuid4
from tempfile import TemporaryDirectory
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.workspace import ImportFormat
from databricks.sdk.service.pipelines import NotebookLibrary, PipelinesEnvironment, PipelineLibrary
from databricks.sdk.service.jobs import NotebookTask, PipelineTask, Task
Expand All @@ -17,8 +16,7 @@
logger = logging.getLogger(__name__)


def test_run_dqx_demo_library(make_notebook, make_schema, make_job, library_ref):
ws = WorkspaceClient()
def test_run_dqx_demo_library(ws, make_notebook, make_schema, make_job, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_library.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand Down Expand Up @@ -47,8 +45,7 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job, library_ref)
logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_library")


def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, make_job, library_ref):
ws = WorkspaceClient()
def test_run_dqx_manufacturing_demo(ws, make_notebook, make_directory, make_schema, make_job, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_manufacturing_demo.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand All @@ -73,8 +70,7 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema,
logging.info(f"Job run {run.run_id} completed successfully for dqx_manufacturing_demo")


def test_run_dqx_quick_start_demo_library(make_notebook, make_job, library_ref):
ws = WorkspaceClient()
def test_run_dqx_quick_start_demo_library(ws, make_notebook, make_job, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_quick_start_demo_library.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand All @@ -92,8 +88,7 @@ def test_run_dqx_quick_start_demo_library(make_notebook, make_job, library_ref):
logging.info(f"Job run {run.run_id} completed successfully for dqx_quick_start_demo_library")


def test_run_dqx_demo_pii_detection(make_notebook, make_job, library_ref):
ws = WorkspaceClient()
def test_run_dqx_demo_pii_detection(ws, make_notebook, make_job, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_pii_detection.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand All @@ -114,14 +109,31 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_job, library_ref):
logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_pii_detection")


def test_run_dqx_dlt_demo(make_notebook, make_pipeline, make_job, library_ref):
ws = WorkspaceClient()
def test_run_dqx_dlt_demo(
skip_if_runtime_not_serverless_compatible, ws, make_notebook, make_schema, make_pipeline, make_job, library_ref
):
"""
Test running the DLT demo notebook in a serverless pipeline.
No need to trigger from non-serverless runtime, since the dlt pipeline use own cluster anyway.
"""
catalog = TEST_CATALOG
schema = make_schema(catalog_name=catalog).name

path = Path(__file__).parent.parent.parent / "demos" / "dqx_dlt_demo.py"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we add more complete DLT tests as a follow-up issue?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can add it as part of this PR, but what does it mean more complete tests? e.g. containing materialized views?

with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)

notebook_path = notebook.as_fuse().as_posix()
pipeline = make_pipeline(
Comment thread
mwojtyczka marked this conversation as resolved.
# DLT / Lakeflow support 3 modes of execution:
# * Full Unity Catalog (UC) mode, the so called DPM (Direct Publishing Mode).
# This mode supports write to arbitrary catalogs and schemas, and performs additional checks, e.g.
# prevent usage of collect() or schema inspection
# * non-UC mode (legacy) where neither catalog & schema nor target are provided
# * UC legacy mode which specifies the target schema where all datasets defined in the pipeline are published
# As part of this test we use the latest full UC mode.
catalog=catalog,
schema=schema,
libraries=[PipelineLibrary(notebook=NotebookLibrary(notebook_path))],
environment=PipelinesEnvironment(dependencies=[library_ref]),
)
Expand All @@ -137,7 +149,7 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline, make_job, library_ref):
logging.info(f"Job run {run.run_id} completed successfully for dqx_dlt_demo")


def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_job):
def test_run_dqx_demo_tool(ws, installation_ctx, make_schema, make_notebook, make_job):
catalog = TEST_CATALOG
schema = make_schema(catalog_name=catalog).name
installation_ctx.replace(
Expand All @@ -152,7 +164,6 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo
install_path = installation_ctx.installation.install_folder()

path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_tool.py"
ws = WorkspaceClient()
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)

Expand All @@ -175,8 +186,7 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo
logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_tool")


def test_run_dqx_streaming_demo_native(make_notebook, make_schema, make_job, tmp_path, library_ref):
ws = WorkspaceClient()
def test_run_dqx_streaming_demo_native(ws, make_notebook, make_schema, make_job, tmp_path, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_streaming_demo_native.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand Down Expand Up @@ -205,8 +215,7 @@ def test_run_dqx_streaming_demo_native(make_notebook, make_schema, make_job, tmp
logging.info(f"Job run {run.run_id} completed successfully for dqx_streaming_demo")


def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path, library_ref):
ws = WorkspaceClient()
def test_run_dqx_streaming_demo_diy(ws, make_notebook, make_job, tmp_path, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_streaming_demo_diy.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand Down Expand Up @@ -263,9 +272,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, make_random, library_ref):
subprocess.run([cli_path, "bundle", "destroy", "--auto-approve"], check=True, capture_output=True, cwd=path)


def test_run_dqx_multi_table_demo(make_notebook, make_schema, make_job, library_ref):

ws = WorkspaceClient()
def test_run_dqx_multi_table_demo(ws, make_notebook, make_schema, make_job, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_multi_table_demo.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand All @@ -288,8 +295,7 @@ def test_run_dqx_multi_table_demo(make_notebook, make_schema, make_job, library_
logging.info(f"Job run {run.run_id} completed successfully for dqx_multi_table_demo")


def test_run_dqx_demo_summary_metrics(make_notebook, make_schema, make_job, library_ref):
ws = WorkspaceClient()
def test_run_dqx_demo_summary_metrics(ws, make_notebook, make_schema, make_job, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_summary_metrics.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand Down Expand Up @@ -318,8 +324,7 @@ def test_run_dqx_demo_summary_metrics(make_notebook, make_schema, make_job, libr
logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_summary_metrics")


def test_run_dqx_ai_assisted_quality_checks_generation(make_notebook, make_job, library_ref):
ws = WorkspaceClient()
def test_run_dqx_ai_assisted_quality_checks_generation(ws, make_notebook, make_job, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_ai_assisted_checks_generation.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand All @@ -337,9 +342,8 @@ def test_run_dqx_ai_assisted_quality_checks_generation(make_notebook, make_job,
logging.info(f"Job run {run.run_id} completed successfully for dqx_quick_start_demo_library")


def test_run_dqx_demo_datacontract_odcs(make_notebook, make_job, library_ref):
def test_run_dqx_demo_datacontract_odcs(ws, make_notebook, make_job, library_ref):
"""Test the ODCS v3.x data contract demo notebook."""
ws = WorkspaceClient()
path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_datacontract_odcs.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand All @@ -357,8 +361,7 @@ def test_run_dqx_demo_datacontract_odcs(make_notebook, make_job, library_ref):
logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_datacontract_odcs")


def test_run_dqx_demo_llm_pk_detection(make_notebook, make_job, library_ref):
ws = WorkspaceClient()
def test_run_dqx_demo_llm_pk_detection(ws, make_notebook, make_job, library_ref):
path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_llm_pk_detection.py"
with open(path, "rb") as f:
notebook = make_notebook(content=f, format=ImportFormat.SOURCE)
Expand Down
Loading