如何unittest pyspark ` `withColumn`‘action - Python 3？

import unittest
from pyspark.sql import SparkSession
from pyspark.sql.functions import col

class SparkUnitTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # 创建SparkSession
        cls.spark = SparkSession.builder \
            .appName("SparkUnitTest") \
            .master("local[*]") \
            .getOrCreate()

    @classmethod
    def tearDownClass(cls):
        # 停止SparkSession
        cls.spark.stop()

    def test_withColumn_action(self):
        # 创建测试数据
        data = [("Alice", 25), ("Bob", 30), ("Charlie", 35)]
        df = self.spark.createDataFrame(data, ["name", "age"])

        # 执行withColumn操作
        df = df.withColumn("age_plus_10", col("age") + 10)

        # 验证结果
        expected_data = [("Alice", 25, 35), ("Bob", 30, 40), ("Charlie", 35, 45)]
        expected_df = self.spark.createDataFrame(expected_data, ["name", "age", "age_plus_10"])
        self.assertEqual(df.collect(), expected_df.collect())

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