Addressed comments.

Author: tdas

streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala

@@ -17,13 +17,12 @@
 
 package org.apache.spark.streaming
 
-import java.io.{NotSerializableException, InputStream}
+import java.io.{InputStream, NotSerializableException}
 import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
 
 import scala.collection.Map
 import scala.collection.mutable.Queue
 import scala.reflect.ClassTag
-import scala.util.control.NonFatal
 
 import akka.actor.{Props, SupervisorStrategy}
 import org.apache.hadoop.conf.Configuration
@@ -43,7 +42,7 @@ import org.apache.spark.streaming.dstream._
 import org.apache.spark.streaming.receiver.{ActorReceiver, ActorSupervisorStrategy, Receiver}
 import org.apache.spark.streaming.scheduler.{JobScheduler, StreamingListener}
 import org.apache.spark.streaming.ui.{StreamingJobProgressListener, StreamingTab}
-import org.apache.spark.util.{Utils, CallSite}
+import org.apache.spark.util.CallSite
 
 /**
  * Main entry point for Spark Streaming functionality. It provides methods used to create
@@ -542,7 +541,7 @@ class StreamingContext private[streaming] (
         Checkpoint.serialize(checkpoint, conf)
       } catch {
         case e: NotSerializableException =>
-          throw new IllegalArgumentException(
+          throw new NotSerializableException(
             "DStream checkpointing has been enabled but the DStreams with their functions " +
               "are not serializable\nSerialization stack:\n" +
               SerializationDebugger.find(checkpoint).map("\t- " + _).mkString("\n")

streaming/src/test/scala/org/apache/spark/streaming/StreamingContextSuite.scala

@@ -17,21 +17,21 @@
 
 package org.apache.spark.streaming
 
-import java.io.File
+import java.io.{File, NotSerializableException}
 import java.util.concurrent.atomic.AtomicInteger
 
 import org.apache.commons.io.FileUtils
-import org.scalatest.{Assertions, BeforeAndAfter, FunSuite}
-import org.scalatest.concurrent.Timeouts
 import org.scalatest.concurrent.Eventually._
+import org.scalatest.concurrent.Timeouts
 import org.scalatest.exceptions.TestFailedDueToTimeoutException
 import org.scalatest.time.SpanSugar._
+import org.scalatest.{Assertions, BeforeAndAfter, FunSuite}
 
-import org.apache.spark.{Logging, SparkConf, SparkContext, SparkException}
 import org.apache.spark.storage.StorageLevel
 import org.apache.spark.streaming.dstream.DStream
 import org.apache.spark.streaming.receiver.Receiver
 import org.apache.spark.util.Utils
+import org.apache.spark.{Logging, SparkConf, SparkContext, SparkException}
 
 
 class StreamingContextSuite extends FunSuite with BeforeAndAfter with Timeouts with Logging {
@@ -143,9 +143,10 @@ class StreamingContextSuite extends FunSuite with BeforeAndAfter with Timeouts w
     }
 
     // Test whether start() fails early when checkpointing is enabled
-    intercept[IllegalArgumentException] {
+    val exception = intercept[NotSerializableException] {
       ssc.start()
     }
+    assert(exception.getMessage().contains("DStreams with their functions are not serializable"))
     assert(ssc.getState() !== StreamingContextState.ACTIVE)
     assert(StreamingContext.getActive().isEmpty)
   }