StreamingSpec.scala

trait StreamingSpec extends AbstractSpec with SharedEmbeddedCassandra with SparkTemplate with BeforeAndAfterAll {
  import org.apache.spark.streaming.StreamingContext
  import scala.concurrent.duration._

  val duration = 10.seconds

  useCassandraConfig(Seq("cassandra-default.yaml.template"))

  def withStreamingContext(testCode: (StreamingContext) => Any, checkpointed: Option[(()=>StreamingContext,String)] = None): Unit = {
    val ssc = checkpointed match {
      case Some((contextGenerator, checkpointDir)) =>
        StreamingContext.getOrCreate(checkpointDir, contextGenerator)
      case None =>
        new StreamingContext(sc, Milliseconds(200))
    }
    try {
      testCode (ssc)
    }
    finally ssc.stop (stopSparkContext = checkpointed.isDefined, stopGracefully = true)
  }


}