我如何测试这个功能？

Question

我有这个功能：

# spark already defined somewhere as:
spark = SparkSession.builder.appName("App").getOrCreate()

def read_data(spark):
    query = "SELECT * FROM table"
    pandas_df = pd.read_sql(query, conn)
    return spark.createDataFrame(pandas_df)

测试一下：

from unittest import mock

@mock.patch("pandas.read_sql")
@mock.patch("pyspark.sql.SparkSession", autospec=True)
def test_read_data(spark_session, pandas_read_sql):
    result = read_data(spark_session)
    assert == ???

我应该以什么方式测试这个是否有意义？任何帮助表示赞赏。

Answer 1

为了测试您的功能，您只需要模拟 pandas.read_sql，不能模拟 spark_session，您需要有一个实例来正确测试您的功能。您可以创建自己的 pytest.fixture 来满足此要求。

from unittest.mock import patch

import pandas
import pyspark.sql
import pytest
from pyspark.sql import SparkSession

from your_module import read_data


@pytest.fixture
def spark_session():
    _spark_session = SparkSession.builder.appName("unit-tests").getOrCreate()
    yield _spark_session
    _spark_session.stop()


@patch("pandas.read_sql")
def test_read_data(mock_read_sql, spark_session):
    # given:
    mock_read_sql.return_value = pandas.DataFrame(
        [(1, "row1"), (2, "row2")], columns=["id", "column1"]
    )

    # when:
    spark_df = read_data(spark_session)

    # then:
    assert isinstance(spark_df, pyspark.sql.DataFrame)

您可以做更多的断言并检查创建的数据框是否具有正确的模式并包含您期望的值。

提示：您应该查看 spark sql 功能，因为您可能不需要使用 pandas 来查询您的数据库。

我如何测试这个功能？

How do I test this function?

python

pandas

python-unittest

pyspark