使用testcontainers测试kafka和spark

<dependency>
    <groupId>org.testcontainers</groupId>
    <artifactId>kafka</artifactId>
    <version>1.16.0</version>
    <scope>test</scope>
</dependency>

<dependency>
    <groupId>org.testcontainers</groupId>
    <artifactId>spark</artifactId>
    <version>1.16.0</version>
    <scope>test</scope>
</dependency>

import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.streaming.StreamingQuery;
import org.apache.spark.sql.streaming.StreamingQueryException;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.KafkaContainer;
import org.testcontainers.containers.SparkContainer;
import org.testcontainers.utility.DockerImageName;

import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeoutException;

import static org.junit.jupiter.api.Assertions.assertEquals;

public class KafkaSparkTest {

    private KafkaContainer kafkaContainer;
    private SparkContainer sparkContainer;

    @BeforeEach
    public void setup() {
        kafkaContainer = new KafkaContainer(DockerImageName.parse("confluentinc/cp-kafka:5.5.1"));
        kafkaContainer.start();

        sparkContainer = new SparkContainer(DockerImageName.parse("apache/spark:3.0.1"));
        sparkContainer.dependsOn(kafkaContainer);
        sparkContainer.start();
    }

    @AfterEach
    public void teardown() {
        sparkContainer.stop();
        kafkaContainer.stop();
    }

    @Test
    public void testKafkaAndSpark() throws TimeoutException, StreamingQueryException {
        // 创建Kafka生产者并发送消息
        kafkaContainer.getKafkaProducer().send(new ProducerRecord<>("test-topic", "key", "value"));

        // 创建SparkSession
        SparkSession sparkSession = SparkSession.builder()
                .appName("KafkaSparkTest")
                .master(sparkContainer.getSparkMasterUrl())
                .getOrCreate();

        // 从Kafka主题中读取数据
        Dataset<Row> kafkaData = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", kafkaContainer.getBootstrapServers())
                .option("subscribe", "test-topic")
                .load();

        // 处理数据
        Dataset<Row> processedData = kafkaData.selectExpr("CAST(value AS STRING)");

        // 启动流处理查询
        StreamingQuery query = processedData.writeStream()
                .format("console")
                .start();

        // 等待查询完成
        query.awaitTermination(5000);

        // 检查输出结果
        List<String> output = sparkContainer.getLogs();
        List<String> expectedOutput = Arrays.asList("value");
        assertEquals(expectedOutput, output);
    }
}