3333
3434from pyspark .context import SparkConf , SparkContext , RDD
3535from pyspark .streaming .context import StreamingContext
36- from pyspark .streaming .kafka import Broker , KafkaUtils , OffsetRange , TopicAndPartition
3736
3837
3938class StreamingTestCase (SparkTestingBaseReuse ):
4039
41- """Basic common test case for Spark Streaming tests. Provides a Spark Streaming context,
42- as well as some helper methods for creating streaming input and collecting
43- streaming output. Based on PySparkStreamingTestCase."""
40+ """Basic common test case for Spark Streaming tests. Provides a
41+ Spark Streaming context as well as some helper methods for creating
42+ streaming input and collecting streaming output.
43+ Modeled after PySparkStreamingTestCase."""
4444
4545 timeout = 10 # seconds
4646 duration = .5
@@ -49,7 +49,6 @@ class StreamingTestCase(SparkTestingBaseReuse):
4949 def setUpClass (cls ):
5050 super (StreamingTestCase , cls ).setUpClass ()
5151 cls .sc .setCheckpointDir ("/tmp" )
52-
5352
5453 @classmethod
5554 def tearDownClass (cls ):
@@ -109,23 +108,25 @@ def get_output(_, rdd):
109108
110109 def run_func (self , input , func , expected , sort = False , input2 = None ):
111110 """
112- @param input: dataset for the test. This should be list of lists or list of RDDs.
113- @param input2: Optional second dataset for the test. If provided func must
114- take two PythonDStreams as input.
115- @param func: wrapped function. This function should return PythonDStream object.
111+ @param input: dataset for the test. This should be list of lists
112+ or list of RDDs.
113+ @param input2: Optional second dataset for the test. If provided your
114+ func must take two PythonDStreams as input.
115+ @param func: wrapped function. This function should return
116+ PythonDStream.
116117 @param expected: expected output for this testcase.
117- Warning: If output is longer than expected this will silently discard the additional
118- output. TODO: fail when this happens.
118+ Warning: If output is longer than expected this will silently
119+ discard the additional output. TODO: fail when this happens.
119120 """
120121 if not isinstance (input [0 ], RDD ):
121122 input = [self .sc .parallelize (d , 1 ) for d in input ]
122123 input_stream = self .ssc .queueStream (input )
123124 if input2 and not isinstance (input2 [0 ], RDD ):
124125 input2 = [self .sc .parallelize (d , 1 ) for d in input2 ]
125- input_stream2 = self .ssc .queueStream (input2 ) if input2 is not None else None
126126
127127 # Apply test function to stream.
128128 if input2 :
129+ input_stream2 = self .ssc .queueStream (input2 )
129130 stream = func (input_stream , input_stream2 )
130131 else :
131132 stream = func (input_stream )
0 commit comments