Package: databricks-test
databricks_testAboutAn experimental unit test framework for Databricks notebooks.This open-source project is not developed by nor affiliated with Databricks.Installingpip install databricks_testUsageAdd a cell at the beginning of your Databricks notebook:# Instrument for unit tests. This is only executed in local unit tests, not in Databricks.if'dbutils'notinlocals():importdatabricks_testdatabricks_test.inject_variables()Theifclause causes the inner code to be skipped when run in Databricks.
Therefore there is no need to install thedatabricks_testmodule on your Databricks environment.Add your notebook into a code project, for example usingGitHub version control in Azure Databricks.Set up pytest in your code project (outside of Databricks).Create a test case with the following structure:importdatabricks_testdeftest_method():withdatabricks_test.session()asdbrickstest:# Set up mocks on dbrickstest# ...# Run notebookdbrickstest.run_notebook("notebook_dir","notebook_name_without_py_suffix")# Test assertions# ...You can set upmocksondbrickstest, for example:dbrickstest.dbutils.widgets.get.return_value="myvalue"See samples below for more examples.Supported featuresSpark context injected into Databricks notebooks:spark,table,sqletc.PySpark with all Spark features including reading and writing to disk, UDFs and Pandas UDFsDatabricks Utilities (dbutils,display) with user-configurable mocksMocking connectors such as Azure Storage, S3 and SQL Data WarehouseUnsupported featuresNotebook formats other than.py(.ipynb,.dbc) are not supportedNon-python cells such as%scalaand%sql(those cells are skipped, as they are stored in.pynotebooks as comments)Writing directly to/dbfsmount on local filesystem:
write to a local temporary file instead and use dbutils.fs.cp() to copy to DBFS, which you can intercept with a mockDatabricks extensions to Spark such asspark.read.format("binaryFile")Sample testSample test case for an ETL notebook reading CSV and writing Parquet.importpandasaspdimportdatabricks_testfromtempfileimportTemporaryDirectoryfrompandas.testingimportassert_frame_equaldeftest_etl():withdatabricks_test.session()asdbrickstest:withTemporaryDirectory()astmp_dir:out_dir=f"{tmp_dir}/out"# Provide input and output location as widgets to notebookswitch={"input":"tests/etl_input.csv","output":out_dir,}dbrickstest.dbutils.widgets.get.side_effect=lambdax:switch.get(x,"")# Run notebookdbrickstest.run_notebook(".","etl_notebook")# Notebook produces a Parquet file (directory)resultDF=pd.read_parquet(out_dir)# Compare produced Parquet file and expected CSV fileexpectedDF=pd.read_csv("tests/etl_expected.csv")assert_frame_equal(expectedDF,resultDF,check_dtype=False)In the notebook, we pass parameters using widgets.
This makes it easy to pass
a local file location in tests, and a remote URL (such as Azure Storage or S3)
in production.# Databricks notebook source# This notebook processed the training dataset (imported by Data Factory)# and computes a cleaned dataset with additional features such as city.frompyspark.sql.typesimportStructType,StructFieldfrompyspark.sql.typesimportDoubleType,IntegerTypefrompyspark.sql.functionsimportcol,pandas_udf,PandasUDFType# COMMAND ----------# Instrument for unit tests. This is only executed in local unit tests, not in Databricks.if'dbutils'notinlocals():importdatabricks_testdatabricks_test.inject_variables()# COMMAND ----------# Widgets for interactive development.dbutils.widgets.text("input","")dbutils.widgets.text("output","")dbutils.widgets.text("secretscope","")dbutils.widgets.text("secretname","")dbutils.widgets.text("keyname","")# COMMAND ----------# Set up storage credentialsspark.conf.set(dbutils.widgets.get("keyname"),dbutils.secrets.get(scope=dbutils.widgets.get("secretscope"),key=dbutils.widgets.get("secretname")),)# COMMAND ----------# Import CSV filesschema=StructType([StructField("aDouble",DoubleType(),nullable=False),StructField("anInteger",IntegerType(),nullable=False),])df=(spark.read.format("csv").options(header="true",mode="FAILFAST").schema(schema).load(dbutils.widgets.get('input')))display(df)# COMMAND ----------df.count()# COMMAND ----------# Inputs and output are pandas.Series of doubles@pandas_udf('integer',PandasUDFType.SCALAR)defsquare(x):returnx*x# COMMAND ----------# Write out Parquet data(df.withColumn("aSquaredInteger",square(col("anInteger"))).write.parquet(dbutils.widgets.get('output')))Advanced mockingSample test case mocking PySpark classes for a notebook connecting to Azure SQL Data Warehouse.importdatabricks_testimportpysparkimportpyspark.sql.functionsasFfromtempfileimportTemporaryDirectoryfrompandas.testingimportassert_frame_equalimportpandasaspddeftest_sqldw(monkeypatch):withdatabricks_test.session()asdbrickstest,TemporaryDirectory()astmp:out_dir=f"{tmp}/out"# Mock SQL DW loader, creating a Spark DataFrame insteaddefmock_load(reader):return(dbrickstest.spark.range(10).withColumn("age",F.col("id")*6).withColumn("salary",F.col("id")*10000))monkeypatch.setattr(pyspark.sql.readwriter.DataFrameReader,"load",mock_load)# Mock SQL DW writer, writing to a local Parquet file insteaddefmock_save(writer):monkeypatch.undo()writer.format("parquet")writer.save(out_dir)monkeypatch.setattr(pyspark.sql.readwriter.DataFrameWriter,"save",mock_save)# Run notebookdbrickstest.run_notebook(".","sqldw_notebook")# Notebook produces a Parquet file (directory)resultDF=pd.read_parquet(out_dir)# Compare produced Parquet file and expected CSV fileexpectedDF=pd.read_csv("tests/sqldw_expected.csv")assert_frame_equal(expectedDF,resultDF,check_dtype=False)IssuesPlease report issues athttps://github.com/microsoft/DataOps/issues.