我们如何使用JUnit进行Spark Dataframe测试？

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;
import static org.junit.Assert.*;

public class DataFrameTest {

    private static SparkSession spark;

    @BeforeClass
    public static void setUp() {
        spark = SparkSession.builder()
                .appName("DataFrameTest")
                .master("local[*]")
                .getOrCreate();
    }

    @AfterClass
    public static void tearDown() {
        if (spark != null) {
            spark.stop();
        }
    }

    @Test
    public void testDataFrameCreation() {
        Dataset<Row> df = spark.createDataFrame(
                java.util.Arrays.asList(
                        RowFactory.create(1, "Alice"),
                        RowFactory.create(2, "Bob")
                ),
                new StructType()
                        .add("id", DataTypes.IntegerType)
                        .add("name", DataTypes.StringType)
        );

        assertEquals(2, df.count());
        assertEquals("Alice", df.filter("id = 1").select("name").first().getString(0));
    }
}