Updated testsuites to work with the slack time of file stream.

2 files changed, 13 insertions, 1 deletions

diff --git a/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala
index 4e25c9566c..0347cc1032 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala
@@ -200,6 +200,9 @@ class CheckpointSuite extends TestSuiteBase {
     val clockProperty = System.getProperty("spark.streaming.clock")
     System.clearProperty("spark.streaming.clock")
 
+    // Disable slack time of file stream when testing with local file system
+    FileInputDStream.disableSlackTime()
+
     // Set up the streaming context and input streams
     val testDir = Files.createTempDir()
     var ssc = new StreamingContext(master, framework, Seconds(1))
@@ -300,6 +303,9 @@ class CheckpointSuite extends TestSuiteBase {
     // Enable manual clock back again for other tests
     if (clockProperty != null)
       System.setProperty("spark.streaming.clock", clockProperty)
+
+    // Restore the default slack time
+    FileInputDStream.restoreSlackTime()
   }
 
 
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala
index 62a9f120b4..e506c954ac 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala
@@ -23,7 +23,7 @@ import akka.actor.IOManager
 import akka.actor.Props
 import akka.util.ByteString
 
-import org.apache.spark.streaming.dstream.{NetworkReceiver, SparkFlumeEvent}
+import org.apache.spark.streaming.dstream.{FileInputDStream, NetworkReceiver, SparkFlumeEvent}
 import java.net.{InetSocketAddress, SocketException, Socket, ServerSocket}
 import java.io.{File, BufferedWriter, OutputStreamWriter}
 import java.util.concurrent.{Executors, TimeUnit, ArrayBlockingQueue}
@@ -152,6 +152,9 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
     // Disable manual clock as FileInputDStream does not work with manual clock
     System.clearProperty("spark.streaming.clock")
 
+    // Disable slack time of file stream when testing with local file system
+    FileInputDStream.disableSlackTime()
+
     // Set up the streaming context and input streams
     val testDir = Files.createTempDir()
     val ssc = new StreamingContext(master, framework, batchDuration)
@@ -196,6 +199,9 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
 
     // Enable manual clock back again for other tests
     System.setProperty("spark.streaming.clock", "org.apache.spark.streaming.util.ManualClock")
+
+    // Restore the default slack time
+    FileInputDStream.restoreSlackTime()
   }