Unit Testing Apache Spark Structured Streaming using MemoryStream
Tags: spark, apachespark, testing, bigdata
Published: 2020-08-10 13:00:33

---

firstPublishedAt: 1597055578925
latestPublishedAt: 1597055578925
slug: unit-testing-apache-spark-structured-streaming-using-memorystream
title: Unit Testing Apache Spark Structured Streaming using MemoryStream

---
Unit testing [Apache Spark](https://spark.apache.org/) Structured Streaming jobs using _MemoryStream_ in a non-trivial task. Sadly enough, official Spark [documentation ](https://spark.apache.org/docs/latest/)still lacks a section on testing. In this post, therefore, I will show you how to start writing unit tests of Spark Structured Streaming.

# What is MemoryStream?

> _A Source that produces value stored in memory as they are added by the user. This Source is intended for use in unit tests as it can only replay data when the object is still available._

> [_Spark SQL Docs_](https://www.javadoc.io/doc/org.apache.spark/spark-sql_2.12/latest/org/apache/spark/sql/execution/streaming/MemoryStream.html)

_MemoryStream_ is one of the streaming sources available in Apache Spark. This source allows us to add and store data in memory, which is very convenient for unit testing. The official docs emphasize this, along with a warning that data can be replayed only when the object is still available.

The _MemoryStream_ takes a type parameter, which gives our unit tests a very desired type safety. The API of this abstraction is rich but not overwhelming. You should check it out for a full reference.

Before writing any unit test, let’s create some sample job which will be tested.

# Writing Spark Structured Streaming job

The job I will be using for the testing, has a simple role — read data from Kafka data source and write it to the Mongo database. In your case, the set of transformations and aggregations will be probably much richer, but the principles stay the same.

First, I read the Kafka data source and extract the `value` column. I am also applying a prebuilt `rsvpStruct` schema, but that is specific to my code sample.

```
import spark.implicits._
val df = spark.readStream
  .format("kafka")
  .option("kafka.bootstrap.servers", "localhost:9092")
  .option("subscribe", "rsvp")
  .load()
val rsvpJsonDf = df.selectExpr("CAST(value as STRING)")
val rsvpStruct = Schema.getRsvpStructSchema
val rsvpDf = rsvpJsonDf.select(from_json($"value", rsvpStruct).as("rsvp"))
```

Next, I am outputting the content of parsed data into a Mongo collection. That is done is `foreachBatch` action, which loops through every batch in the current result set and allows to perform an arbitrary operation. In my case, I am using [Mongo Spark](https://docs.mongodb.com/spark-connector/current/scala-api/) connector for persisting the batches.

```
rsvpDf.writeStream
  .outputMode("append")
  .option("checkpointLocation", "/raw")
  .foreachBatch({(batchDf: DataFrame, batchId: Long) =>
    val outputDf = batchDf.select("rsvp.*")
    outputDf.write
      .format("mongo")
      .mode("append")
      .save()
  })
  .start()
  .awaitTermination()
```

Now that we have the sample job created, let’s start writing a unit test using _MemoryStream_ class.

# Unit testing Spark using MemoryStream

The preliminary step in writing a unit test is to have some sample data that can be used by unit tests. In my case, I am using a stringified JSON, that contains information I am normally receiving from my Kafka data source. You should use any data that fits your use case.

```
val sampleRsvp = """{"venue":{"venue_name":"Capitello Wines" ...
```

Next, inside the test case, we create our `SparkSession`. Nothing special has to be set here, just typical config you use.

```
class SaveRawDataTest extends AnyFunSuite {
  test("should ingest raw RSVP data and read the parsed JSON correctly") {
    val spark = SparkSession.builder()
      .appName("Save Raw Data Test")
      .master("local[1]")
      .getOrCreate()
...
```

Going further, we define our `MemoryStream` using String type parameter. I use string, because that's the type of sample data I am using. To create the `MemoryStream`, to imports are needed: `sqlContext` that can be obtained from SparkSession and Spark `implicits` library, which includes the needed type encoders. Lastly, I convert the stream to `DataSet`.

```
implicit val sqlCtx: SQLContext = spark.sqlContext
import spark.implicits._

val events = MemoryStream[String]
val sessions = events.toDS
```

Converting to `DataSet` allows us to make first assertion. We can check if the `DataSet` is indeed a streaming one. You can do this using the following code.

```
assert(sessions.isStreaming, "sessions must be a streaming Dataset")
```

Next, I add the transformations that have been included in the sample job. In this case however, I will not be checking if the data has been saved in Mongo — I am only interested if the raw stringified JSON has been parse correctly and I can run SQL queries on it. The transformations I use are:

```
val transformedSessions = sessions.select(from_json($"value", rsvpStruct).as("rsvp"))
val streamingQuery = transformedSessions
  .writeStream
  .format("memory")
  .queryName("rawRsvp")
  .outputMode("append")
  .start
```

It is important to use `memory` format, so that the actual data can be queried in a further step to make the assertions. It is useful to also name those results, using `queryName` option.

Lastly, I add my sample data into the instance of `MemoryStream`, process those and commit the offsets, as shown below:

```
val batch = sampleRsvp
val currentOffset = events.addData(batch)

streamingQuery.processAllAvailable()
events.commit(currentOffset.asInstanceOf[LongOffset])
```

The very last step is to query the committed data and run some assertions on them. In my case, I run a `SQL` query, to get the event_name property and I run assertions against that.

```
val rsvpEventName = spark.sql("select rsvp.event.event_name from rawRsvp")
    .collect()
    .map(_.getAs[String](0))
    .head

assert(rsvpEventName == "Eugene ISSA & Technology Association of Oregon December Cyber Security Meetup")
```

# Summary

I hope you have found this post useful. If so, don’t hesitate to like or share this post. Additionally, you can follow me on my social media if you fancy so 🙂