我如何測試這個功能？

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 來查詢您的數據庫。