Unit Testing Functions in Databricks Notebooks

1.What a unit test checks, and how to write a function worth testing

On Databricks, unit testing means testing self-contained pieces of code, usually individual functions, early and often. The aim is to find problems and mistaken assumptions sooner. The Databricks guide on this topic sets a clear boundary. It covers unit tests for functions and says it does not cover integration, system, acceptance, or performance testing. Keep that boundary in mind for the exam. A unit test calls one function with specific inputs and asserts on what comes back. Checking whether whole pipelines work together is a separate kind of test.

The design rule is simple. To get the best results, a function should return one predictable outcome of one data type. An existence check should return a boolean. A row count should return a non-negative whole number. The docs name the anti-pattern directly: a function that returns false when something is missing but returns the thing itself when it exists. A count function that returns false when there are no rows is the same mistake. The example functions below follow the rule. They live in a plain myfunctions.py file in a Git folder rather than in a notebook. Because of that, the file has to build its own Spark session, which a notebook would normally create for you.

# Does the specified column exist in the given DataFrame?
def columnExists(dataFrame, columnName):
  if columnName in dataFrame.columns:
    return True
  else:
    return False

# How many rows are there for the specified value in the specified column
# in the given DataFrame?
def numRowsInColumnForValue(dataFrame, columnName, columnValue):
  df = dataFrame.filter(col(columnName) == columnValue)

  return df.count()

Unit tests are one of two kinds of test in the Databricks MLOps workflow. The other is the integration test, and the two run at different points. In the development stage, the focus is experimentation, and you write the code and its unit tests. The staging stage is where the ML pipeline code is tested to make sure it is ready for production. ML engineers build a CI pipeline for it, and a pull request triggers it. The CI pipeline first builds the source code and runs unit tests. If they fail, the pull request is rejected. Then it runs integration tests. An integration test runs all the pipelines together, including feature engineering, model training, inference, and monitoring, and checks that they function correctly together. That is the difference from a unit test, which checks one function in isolation. The staging environment should match production as closely as is reasonable, and it has its own catalog in Unity Catalog. Assets written there are generally temporary. Integration tests can trade fidelity for speed or cost, for example by training on small subsets of data. They can also be scheduled periodically on the main branch. The production stage does not repeat that testing. The release branch only starts the production stage after all CI tests have passed. Keep unit tests away from production data, because they should run against non-production data, for example views instead of tables.

Sources123

2.Writing and running unit tests in a notebook

Databricks recommends writing and running unit tests in a notebook. You can run some commands in the web terminal, but it has more limitations, including no Spark support. Many unit testing libraries work directly in a notebook cell. The built-in Python unittest package is one example. In the sample below, the test class inherits from unittest.TestCase and the test method compares a known input with a known output. The unittest.main(argv=[''], verbosity=2, exit=False) call runs the tests inside the cell. The final assert on wasSuccessful() makes the cell itself fail when any test fails.

def reverse(s):
    return s[::-1]

import unittest

class TestHelpers(unittest.TestCase):
    def test_reverse(self):
        self.assertEqual(reverse('abc'), 'cba')

r = unittest.main(argv=[''], verbosity=2, exit=False)
assert r.result.wasSuccessful(), 'Test failed; see logs above'

Test failures show up in the cell's output area. You can hide test code or results from the cell actions menu, and errors are still displayed even when results are hidden. A widget lets one notebook switch between running its tests and doing its normal work. A Mode dropdown chooses the branch.

dbutils.widgets.dropdown("Mode", "Test", ["Test", "Normal"])

def reverse(s):
  return s[::-1]

if dbutils.widgets. ? ('Mode') == 'Test':
  assert reverse('abc') == 'cba'
  print('Tests passed')
else:
  print(reverse('desrever'))

For tests stored in Python files, the workspace has a testing sidebar, inline run controls, and a results pane. Discovery uses pytest naming conventions. A file counts as a test file only if its name matches test_*.py or *_test.py. Inside a test file, a test is a test-prefixed function outside a class, or a test-prefixed method inside a Test-prefixed class that has no __init__ method. Static and class methods inside Test classes are discovered too. A function called verify_output is never collected, however sensible its name looks.

class TestClass():
    def test_1(self):
        assert True

    def test_3(self):
        assert 4 == 3

def test_foo():
    assert "foo" == "bar"

You can also run tests without anyone starting them. A scheduled notebook job can run them on a regular basis and email the address you specify. For notebooks in a Databricks Git folder, you can set up a CI/CD-style workflow in which the notebook tests run on every commit, for example with Databricks GitHub Actions.

Sources145

3.Organizing functions and their unit tests

Where you put a function and where you put its tests both matter. That choice decides who can reuse the function and which tools can test it. For Python, R, and Scala notebooks, Databricks describes three common layouts, each with its own benefits and challenges. SQL notebooks follow a separate pattern.

The separate-notebook layout depends on %run. The test notebook's first cell runs %run ./shared-code-notebook, which loads reverse into the test notebook. A later cell then holds the unittest class. Databricks Git folders give you the file-based alternative: test code can live in non-notebook source files and still be called from a notebook.

Sources14

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.pytest will run any function in a test file, whatever its name, such as verify_nulls or check_output.Why is that wrong?

   Discovery follows pytest naming conventions. Test files must match test_*.py or *_test.py, and only test-prefixed functions, or test-prefixed methods in Test-prefixed classes without __init__, are collected.

   Covered in Writing and running unit tests in a notebook

2. 2.Storing functions and their tests in files outside notebooks is the best practice for every notebook language.Why is that wrong?

   That approach is not supported for Scala notebooks. For Scala, Databricks recommends functions in one notebook and their tests in a separate notebook.

   Covered in Organizing functions and their unit tests

3. 3.Unit tests are the right way to check that feature engineering, training, and inference work together, and they should run against production data.Why is that wrong?

   A unit test checks one function. Integration tests run the pipelines together in staging. Unit tests should run against non-production data, not production data.

   Covered in What a unit test checks, and how to write a function worth testing