Scala Spark Kafka代码函数式改造及单元测试咨询

Great question! Let's break this down into functional refactoring and unit testing best practices tailored to your Scala/Spark/Kafka pipeline.

Functional Style Refactoring #

The key principles here are separating pure logic from side effects, using immutable data structures, and favoring stateless functions over stateful classes. Let's walk through the changes:

1. Define Immutable Data Models #

Start with case classes to represent your domain data—they're immutable by default and play nicely with functional patterns:

import org.json4s.DefaultFormats
import org.json4s.jackson.Serialization.write

case class SubCategory(code: String, name: String)
case class Category(code: String, name: String, subCategories: List[SubCategory])

2. Extract Pure Logic into Stateless Functions #

Move all side-effect-free logic (data transformation, JSON serialization, key generation) into a singleton object. This avoids unnecessary class instances and makes logic reusable/testable:

import org.apache.spark.sql.DataFrame
import java.security.MessageDigest

object CategoryTransformations {
  implicit val formats = DefaultFormats

  def transformToCategories(df: DataFrame): List[Category] = {
    df.collect()
      .groupBy(row => (row.getAs[String]("CAT_CODE"), row.getAs[String]("CAT_NAME")))
      .map { case ((catCode, catName), rows) =>
        val subCats = rows.map(row => 
          SubCategory(row.getAs[String]("SUBCAT_CODE"), row.getAs[String]("SUBCAT_NAME"))
        ).toList
        Category(catCode, catName, subCats)
      }.toList
  }

  def toJson(categories: List[Category]): String = write(categories)

  def generateKafkaKey(json: String): String = {
    MessageDigest.getInstance("SHA-1")
      .digest(json.getBytes("UTF-8"))
      .map("%02x".format(_))
      .mkString
  }
}

3. Isolate Side Effects #

Keep all I/O operations (Spark initialization, Kafka sending) in a separate module. This makes side effects explicit and easier to manage:

import org.apache.spark.sql.SparkSession
import org.slf4j.LoggerFactory

object ExtractProcessor {
  private val Logger = LoggerFactory.getLogger(getClass)

  def run(): Unit = {
    val spark = SparkSession.builder().appName("test app").getOrCreate()
    try {
      val df = spark.sql("""
        SELECT DISTINCT SUBCAT_CODE, SUBCAT_NAME, CAT_CODE, CAT_NAME 
        FROM CATEGORY_HIERARCHY 
        ORDER BY CAT_CODE, SUBCAT_CODE
      """)
      
      val categories = CategoryTransformations.transformToCategories(df)
      val jsonMessage = CategoryTransformations.toJson(categories)
      val kafkaKey = CategoryTransformations.generateKafkaKey(jsonMessage)

      Logger.info(s"Json Message: ${jsonMessage}")
      Logger.info(s"Kafka Key: ${kafkaKey}")
      
      KafkaUtil.send(kafkaKey, jsonMessage, "testTopic")
    } finally {
      spark.stop()
    }
  }
}

4. Simplify Kafka Util #

Refactor KafkaUtil to be stateless (no need for a class—use a singleton object instead):

import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
import java.util.Properties
import org.slf4j.LoggerFactory

object KafkaUtil {
  private val Logger = LoggerFactory.getLogger(getClass)

  private def getProducerProps: Properties = {
    val props = new Properties()
    props.put("bootstrap.servers", "localhost:9092")
    props.put("client.id", "test publisher")
    props.put("key.serializer", "org.apache.kafka.common.serialization.StringSerializer")
    props.put("value.serializer", "org.apache.kafka.common.serialization.StringSerializer")
    props
  }

  def send(key: String, message: String, topicName: String): Unit = {
    val producer = new KafkaProducer[String, String](getProducerProps)
    try {
      val record = new ProducerRecord[String, String](topicName, key, message)
      producer.send(record)
    } finally {
      producer.close()
      Logger.info("Kafka producer closed...")
    }
  }
}

Should You Create a Class with getCategories? #

No—for functional style, stateless functions in a singleton object are preferred. Classes introduce state (even if unintended) and make logic harder to reuse. If you need to vary behavior (e.g., different SQL queries), pass configuration as parameters to the functions instead of using class members.

Functional-Style Unit Testing #

The goal is to test pure logic in isolation and validate side effects with mocks. Use ScalaTest (or Specs2) and Mockito for this.

1. Test Pure Logic (No Side Effects) #

These tests are fast and reliable—no need for Spark or Kafka dependencies:

import org.scalatest.flatspec.AnyFlatSpec
import org.scalatest.matchers.should.Matchers
import org.apache.spark.sql.Row

class CategoryTransformationsSpec extends AnyFlatSpec with Matchers {
  "transformToCategories" should "map DataFrame rows to Category objects correctly" in {
    val testRows = Seq(
      Row("S1", "Subcat 1", "C1", "Cat 1"),
      Row("S2", "Subcat 2", "C1", "Cat 1"),
      Row("S3", "Subcat 3", "C2", "Cat 2")
    )

    val grouped = testRows.groupBy(row => (row(2).toString, row(3).toString))
    val categories = grouped.map { case ((catCode, catName), rows) =>
      val subCats = rows.map(r => SubCategory(r(0).toString, r(1).toString)).toList
      Category(catCode, catName, subCats)
    }.toList

    categories should contain allOf(
      Category("C1", "Cat 1", List(SubCategory("S1", "Subcat 1"), SubCategory("S2", "Subcat 2"))),
      Category("C2", "Cat 2", List(SubCategory("S3", "Subcat 3")))
    )
  }

  "toJson" should "serialize Category list to valid JSON" in {
    val testCategories = List(Category("C1", "Cat 1", List(SubCategory("S1", "Subcat 1"))))
    val json = CategoryTransformations.toJson(testCategories)
    
    json should include(""""code":"C1"""")
    json should include(""""subCategories"""")
    json should include(""""code":"S1"""")
  }

  "generateKafkaKey" should "produce a valid SHA-1 hash" in {
    val testJson = """[{"code":"C1","name":"Cat 1","subCategories":[{"code":"S1","name":"Subcat 1"}]}]"""
    val key = CategoryTransformations.generateKafkaKey(testJson)
    
    key should have length 40
    key should fullyMatch regex "[a-f0-9]{40}"
  }
}

2. Test Side Effects (With Mocks) #

For code that interacts with Spark or Kafka, use mocks to avoid real external calls:

import org.scalatest.flatspec.AnyFlatSpec
import org.scalatest.matchers.should.Matchers
import org.mockito.Mockito._
import org.apache.spark.sql.{SparkSession, DataFrame}

class ExtractProcessorSpec extends AnyFlatSpec with Matchers {
  it should "execute the pipeline and trigger Kafka send" in {
    // Mock Spark components
    val mockSpark = mock(classOf[SparkSession])
    val mockDf = mock(classOf[DataFrame])
    when(mockSpark.sql(anyString())).thenReturn(mockDf)
    
    // Mock Kafka send behavior
    var sendWasCalled = false
    val mockSend = (key: String, msg: String, topic: String) => {
      sendWasCalled = true
      topic shouldBe "testTopic"
      key should not be empty
      msg should not be empty
    }

    // To make this work, modify ExtractProcessor.run to accept a send function as a parameter:
    // def run(sendFunc: (String, String, String) => Unit = KafkaUtil.send): Unit = { ... }
    // Then call ExtractProcessor.run(mockSend) here

    sendWasCalled shouldBe true
  }
}

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