Pyspark unit test: how to mock sql calls (and only the sql calls)?

user_2533340 0
Viewed 931

I'm having trouble testing the following function

# module1.py
from pyspark.sql import DataFrame as SparkDataFrame
from pyspark.sql import SparkSession
from pyspark.sql import functions as f


spark = SparkSession.getActiveSession()


def filter_big_table(
    conditions: Optional[List[str]] = None,
) -> SparkDataFrame:

    sdf = spark.sql(f"select * from SUPER_BIG_TABLE")

    if conditions:
        sdf = sdf.filter(f.col("condition_col").isin(conditions))
    return sdf

Here's how I've been trying to approach it with unittest.mock and pytest:

#conftest.py
import pytest
from pyspark.sql import SparkSession

@pytest.fixture(scope="session")
def spark():
    return SparkSession.builder.master("local").appName("unit_tests").getOrCreate()

And the actual test, where I create a fake DataFrame with tuple fixtures so that I don't have to actually run this query spark.sql(f"select * from SUPER_BIG_TABLE").

# test.py
from unittest.mock import patch

from module1 import filter_big_table
from . import assert_sdf_equal  # Helper function


DATA_FIXTURE = List[Tuple[str, int]]
DATA_SCHEMA = StructType

@patch("pyspark.sql.SparkSession.sql")
def test_filter_big_table(mock_data, spark):
    
    input_sdf = spark.createDataFrame(DATA_FIXTURE, DATA_SCHEMA)
    mock_data.return_value = input_sdf

    expected_sdf = spark.createDataFrame(DATA_FIXTURE, DATA_SCHEMA)
    output_sdf = filter_big_table()

    # Passing no conditions should return the same spark dataframe
    assert_sdf_equal(output_sdf, expected_sdf)

The error I keep getting, however, is AttributeError: 'NoneType' object has no attribute 'sql'.

I feel this is because of the way the module is getting the current spark session on the module1.py, but I didn't want the module to change for the test to pass. I've been toying around with different kinds of mocks, to no avail.

Edit

Tried using @patch("module1.spark") instead, but instead of an AttributeError, I receive a mock object instead of the input_sdf fixture. Also tried replacing mock_data.return_value for mock_data.side_effects, but it didn't help.

pyspark unit-testing python
1 Answers

Had to change the module in order for it to work :(

Now the session is created inside the function.

# module1.py
from pyspark.sql import DataFrame as SparkDataFrame
from pyspark.sql import SparkSession
from pyspark.sql import functions as f


def filter_big_table(
    conditions: Optional[List[str]] = None,
) -> SparkDataFrame:
    spark = SparkSession.getActiveSession()
    sdf = spark.sql(f"select * from SUPER_BIG_TABLE")

    if conditions:
        sdf = sdf.filter(f.col("condition_col").isin(conditions))
    return sdf

And the decorator is @patch("module1.SparkSession.sql")

That way no name collisions occur and the mock successfully replaces the sql method.

user_25333400
Related