topic Re: Python mocking dbutils in unittests in Data Engineering

Python mocking dbutils in unittests

confused_dev — Mon, 31 Oct 2022 19:43:52 GMT

I am trying to write some unittests using pytest, but I am coming accross the problem of how to mock my dbutils method when dbutils isn't being defined in my notebook.

Is there a way to do this so that I can unit test individual functions that are utilizing dbutils?

Re: Python mocking dbutils in unittests

fermin_vicente — Mon, 21 Nov 2022 15:30:49 GMT

Hi,

You can mock dbutils. An easy way to do it is to:

Try to receive dbutils as a parameter in your functions (inject it) instead of using it globally. This way your code is more testable and you won't to do patching which is a bit more cumbersome.
Use a mock library. Unittest.mock is the simplest approach
Call your function passing down a mock instead of the actual dbutils

Example:

your library under test

def my_function(dbutils):
   ...
   dbutils.fs.ls("/tmp")  # this is using the local variable received by parameter
   ...

your notebook

from my_library import my_function
 
...
my_function(dbutils)  # this refers to the global dbutils variable

your test

from unittest.mock import MagicMock
 
from my_library import my_function
 
 
def test_my_function_calls_dbutils():
   mock_dbutils = MagicMock()
 
   my_function(dbutils=mock_dbutils)
 
   mock_dbutils.fs.ls.assert_called_once_with("/tmp")

Re: Python mocking dbutils in unittests

xiangzhu — Mon, 21 Nov 2022 15:40:11 GMT

above reponse is a pure mock, and below is another example from dbx about a fixture with local filesystem, you can also add mock on dbutils.secrets:

https://github.com/databrickslabs/dbx/blob/b2989213b5f67e2b7ccf8adeba97da70e88ffff2/dbx/templates/projects/python_basic/render/%7B%7Bcookiecutter.project_name%7D%7D/tests/unit/conftest.py#L32

Re: Python mocking dbutils in unittests

Anonymous — Mon, 16 Jan 2023 06:13:34 GMT

Hi @Jake P

Hope all is well! Just wanted to check in if you were able to resolve your issue and would you be happy to share the solution or mark an answer as best? Else please let us know if you need more help.

We'd love to hear from you.

Thanks!

Re: Python mocking dbutils in unittests

drorata — Thu, 24 Oct 2024 11:36:03 GMT

How can I add locally type annotation to the example provided above by @fermin_vicente?

The type of `dbutils` seems to be `dbruntime.dbutils.DBUtils` but I'm not sure how to get it available in local dev.

Re: Python mocking dbutils in unittests

fermin_vicente — Fri, 25 Oct 2024 05:57:23 GMT

I personally went without type hint for the dbutils param to avoid extra dependencies just for that, but I think you can actually get that type from the Databricks SDK, according to the docs:

https://docs.databricks.com/en/dev-tools/databricks-connect/python/databricks-utilities.html

https://pypi.org/project/databricks-sdk/#description

I would worry in this case about matching the proper SDK with the DBR you're using to ensure the type validations match (although I guess dbutils doesn't evolve a lot)

Hope this helps

Re: Python mocking dbutils in unittests

saurabh18cs — Fri, 25 Oct 2024 06:56:52 GMT

import unittest
from unittest.mock import MagicMock, patch

# Import the function to be tested
from your_notebook import my_function

class TestMyFunction(unittest.TestCase):
@patch('your_notebook.dbutils')
def test_my_function(self, mock_dbutils):
# Create a mock for dbutils.fs
mock_fs = MagicMock()
mock_dbutils.fs = mock_fs

# Define the behavior of the mock methods
mock_fs.mkdirs.return_value = None
mock_fs.ls.return_value = ["file1", "file2"]

# Call the function to be tested
result = my_function()

# Assertions

if __name__ == '__main__':
unittest.main()

Re: Python mocking dbutils in unittests

pavlosskev — Fri, 25 Oct 2024 09:44:06 GMT

Fermin_vicente's answer is pretty good already. Below is how you can do something similar with conftest.py

# conftest.py import pytest from unittest.mock import MagicMock from pyspark.sql import SparkSession @pytest.fixture(scope="session") def dbutils(): dbutils = MagicMock() yield dbutils dbutils.stop() # Also for the SparkSession @pytest.fixture(scope="session") def spark(): spark = ( SparkSession.builder.master("local[*]").appName("app-name").getOrCreate() ) yield spark spark.stop()

Then for the testing code:

def test_query_generation(spark, dbutils): """ Test example """ test_obj = QueryGenerator(test_config, spark, dbutils) expected_query = ( "SELECT * from community.databricks.com" ) assert ( test_obj.get_query() == expected_query )

Your class object has to get input a dbutils object. I think the same can be done if you want to use just functions.

Re: Python mocking dbutils in unittests

kenmyers-8451 — Tue, 23 Dec 2025 20:01:16 GMT

If this helps anyone here is how we do this:

We rely on databricks_test for injecting dbutils into the notebooks that we're testing (which is a 3rd party package mind you and hasn't been updated in a while but still works). And in our notebooks we put this as the first code cell:

import os # for databricks environment only, dbutils is not immediately available. if "DATABRICKS_RUNTIME_VERSION" in os.environ: print("Found databricks runtime, importing dbutils for etl.utils.load_table ...") from databricks.sdk.runtime import dbutils # Instrument for unit tests. This is intended to only be executed in local unit tests, not in Databricks. if "dbutils" not in locals() and "DATABRICKS_RUNTIME_VERSION" not in os.environ: print("Unable to locate dbutils") import databricks_test databricks_test.inject_variables()

The first conditional is mainly for when you're running this on databricks and probably not necessary since dbutils should be available anyway, but I think it helped with our IDE's understanding what dbutils was and stopped throwing flake errors.

The second conditional is for injecting the databricks_test mocks into the notebook when a test is running.