SparkSuite.scala

//> using scala "2.12.17"
//> using dep "org.apache.spark::spark-sql:3.3.2"
//> using dep "org.scalameta::munit:0.7.29"


import org.apache.spark.sql.{Dataset, SparkSession}
import org.apache.logging.log4j.LogManager
import org.apache.logging.log4j.core.LoggerContext
import org.apache.logging.log4j.Level

import java.io.ByteArrayOutputStream

trait SparkSuite extends munit.FunSuite {

  // turn off noisy logging
  private val _ = {
    val ctx          = LogManager.getContext(false).asInstanceOf[LoggerContext]
    val config       = ctx.getConfiguration()
    val loggerConfig = config.getLoggerConfig(LogManager.ROOT_LOGGER_NAME)
    loggerConfig.setLevel(Level.OFF)
    ctx.updateLoggers()
  }

  implicit lazy val spark =
    SparkSession
      .builder()
      .master("local")
      .appName("test")
      .config("spark.sql.shuffle.partitions", "1")
      .getOrCreate()

  def captureShow(ds: Dataset[_]): String = {
    val outCapture = new ByteArrayOutputStream
    Console.withOut(outCapture)(ds.show())
    new String(outCapture.toByteArray).strip()
  }
}


class ExampleSuite extends SparkSuite {
  import spark.implicits._

  test("example test") {
    val ds = Seq(
      Isb(1, "hello", false),
      Isb(2, "world", true)
    ).toDS()

    val aliased = ds.selectExpr(
      "i as integer",
      "s as string",
      "b as boolean"
    )

    val show = captureShow(aliased)

    val expected =
      """|+-------+------+-------+
         ||integer|string|boolean|
         |+-------+------+-------+
         ||      1| hello|  false|
         ||      2| world|   true|
         |+-------+------+-------+""".stripMargin

    assertEquals(show, expected)
  }
}

private case class Isb(i: Int, s: String, b: Boolean)