pavlosskev
New Contributor III

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.