PySpark脚本函数单元测试：如何提交编写好的测试用例？

Got it, let's walk through how to execute your PySpark unit tests properly. We'll cover both testing your standalone rename_chars function (no Spark needed) and the column_names function (requires a Spark DataFrame). I'll use pytest (the most popular Python testing framework) first, then also show the built-in unittest approach for flexibility.

Step 1: Install Required Dependencies #

First, make sure you have the necessary packages installed:

pip install pytest pyspark

Step 2: Write Your Test File #

Create a test file (e.g., test_spark_transforms.py) with the following code. I'll include your original functions, add test cases, and set up a reusable SparkSession for DataFrame tests:

from pyspark.sql import SparkSession
from functools import reduce
import pytest

# Your original transformation functions
def rename_chars(column_name):
    chars = ((' ', '_&'), ('.', '_$'))
    new_cols = reduce(lambda a, kv: a.replace(*kv), chars, column_name)
    return new_cols

def column_names(df):
    changed_col_names = df.schema.names
    for cols in changed_col_names:
        df = df.withColumnRenamed(cols, rename_chars(cols))
    return df

# Reusable pytest fixture to create/clean up SparkSession
@pytest.fixture(scope="session")
def spark():
    spark = SparkSession.builder \
        .master("local[1]") \
        .appName("PySparkTestSuite") \
        .getOrCreate()
    yield spark
    spark.stop()

# Tests for the standalone rename_chars function
def test_rename_chars_replaces_spaces():
    assert rename_chars("Customer Name") == "Customer_&_Name"

def test_rename_chars_replaces_dots():
    assert rename_chars("Customer.Email") == "Customer_$_Email"

def test_rename_chars_handles_mixed_chars():
    assert rename_chars("Customer Name.Account ID") == "Customer_&_Name_$_Account_&_ID"

# Test for the column_names DataFrame transformation
def test_column_names_updates_all_columns(spark):
    # Create a test DataFrame with original column names
    test_data = [("Alice", 30, "alice@co.com")]
    original_df = spark.createDataFrame(test_data, ["User Name", "Age", "User.Email"])
    
    # Apply your column renaming function
    transformed_df = column_names(original_df)
    
    # Verify the new column names match expectations
    expected_columns = ["User_&_Name", "Age", "User_$_Email"]
    assert transformed_df.schema.names == expected_columns

Step 3: Run the Tests #

Execute the tests with this terminal command:

pytest test_spark_transforms.py -v

Alternative: Using Python's Built-in unittest #

If you prefer sticking to Python's standard library, here's how to structure the test file:

from pyspark.sql import SparkSession
from functools import reduce
import unittest

# Your original functions (same as above)
def rename_chars(column_name):
    chars = ((' ', '_&'), ('.', '_$'))
    new_cols = reduce(lambda a, kv: a.replace(*kv), chars, column_name)
    return new_cols

def column_names(df):
    changed_col_names = df.schema.names
    for cols in changed_col_names:
        df = df.withColumnRenamed(cols, rename_chars(cols))
    return df

class TestSparkColumnTransforms(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Create SparkSession once for all tests
        cls.spark = SparkSession.builder \
            .master("local[1]") \
            .appName("PySparkUnitTests") \
            .getOrCreate()
    
    @classmethod
    def tearDownClass(cls):
        # Clean up SparkSession after tests
        cls.spark.stop()
    
    def test_rename_chars_spaces(self):
        self.assertEqual(rename_chars("Order Date"), "Order_&_Date")
    
    def test_rename_chars_dots(self):
        self.assertEqual(rename_chars("Order.Total"), "Order_$_Total")
    
    def test_column_names_transforms_df(self):
        test_data = [("123", 49.99, "2024-01-01")]
        original_df = self.spark.createDataFrame(test_data, ["Order ID", "Order.Total", "Order Date"])
        transformed_df = column_names(original_df)
        expected_columns = ["Order_&_ID", "Order_$_Total", "Order_&_Date"]
        self.assertEqual(transformed_df.schema.names, expected_columns)

if __name__ == '__main__':
    unittest.main()

Run this version with:

python test_spark_transforms.py

Quick Tips #

• Using local[1] for the Spark master ensures tests run in a lightweight single-node mode, ideal for unit testing.
• The rename_chars function is pure (no side effects), so its tests are simple and fast—no Spark required.
• Always clean up the SparkSession after tests to avoid leaving unnecessary processes running.

内容的提问来源于stack exchange，提问作者User12345