kenmyers-8451
Contributor II

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.