pyspark unit test based on python unittest library
pyspark unit test
Pyspark is a powerful framework for large scale data analysis. Because of the easy-to-use API, you can easily develop pyspark programs if you are familiar with Python programming.
One problem is that it is a little hard to do unit test for pyspark. After some google search using “pyspark unit test”, I only get articles about using py.test or some other complicated libraries for pyspark unit test. However, I don’t want to install any other third party libraries . What I want is to set up the pyspark unit test environment just based on the unittest library, which is currently used by the project.
Fortunately, I found a file from the spark github repository. Based on the code, I made a simple example here to describe the process to setup pyspark unit test environment. The advantage of this method is that the setup is extremely easy comparing with other third party library based Pyspark unit test.
The pyspark unit test base class
There are two base classes defined for pyspark unit test. Both of tem extend the unittest.TestCase class. The first class is the ReusedPySparkTestCase, which can reuse the sparkContext across all unit test methods, as the sparkContext sc is initialized in the setUpClass() method and stopped in the tearDownClass() method.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
class ReusedPySparkTestCase(unittest.TestCase): @classmethod def setUpClass(cls): conf = SparkConf().setMaster("local[2]") \ .setAppName(cls.__name__) \ .set("spark.authenticate.secret", "111111") cls.sc = SparkContext(conf=conf) @classmethod def tearDownClass(cls): cls.sc.stop() class PySparkTestCase(unittest.TestCase): def setUp(self): self._old_sys_path = list(sys.path) conf = SparkConf().setMaster("local[2]") \ .setAppName(self.__class__.__name__) \ .set("spark.authenticate.secret", "111111") self.sc = SparkContext(conf=conf) def tearDown(self): self.sc.stop() sys.path = self._old_sys_path |
Based on the python documentation:
|
1 2 |
setUpClass: is a class method called before tests in an individual class run, tearDownClass(): is class method called after tests in an individual class have run. |
This means setUpClass and tearDownClass are run once for the whole class, so we can share the initialized sparkContext across the test methods.
For the PySparkTestCase, the sparkContext is initialized in the setUp method and stopped in the tearDown method, So each test method will have its own sparkContext. This is because setUp and tearDown are run before and after each test method.
The advantage of using ReusedPySparkTestCase class is that all the unit test methods in the same test class can share or reuse the same sparkContext. If you have many test methods, by reusing the sparkContext can save time as the initialization of the sparkContext is time consuming.
In the following code, I use simple examples to show that all the test methods share the same sparkContext when we extend the ReusedPySparkTestCase class.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import os import sys import unittest SPARK_HOME = os.environ["SPARK_HOME"] os.path.join(SPARK_HOME) print SPARK_HOME # Add the PySpark directories to the Python path: sys.path.insert(1, os.path.join(SPARK_HOME, 'python')) sys.path.insert(1, os.path.join(SPARK_HOME, 'python', 'pyspark')) sys.path.insert(1, os.path.join(SPARK_HOME, 'python', 'build')) sys.path.insert(1, os.path.join(SPARK_HOME, 'python', 'lib/py4j-0.8.2.1-src.zip')) # If PySpark isn't specified, use currently running Python binary: pyspark_python = sys.executable os.environ['PYSPARK_PYTHON'] = pyspark_python from pyspark.conf import SparkConf from pyspark.context import SparkContext sc_values = {} class ReusedPySparkTestCase(unittest.TestCase): @classmethod def setUpClass(cls): conf = SparkConf().setMaster("local[2]") \ .setAppName(cls.__name__) \ .set("spark.authenticate.secret", "111111") cls.sc = SparkContext(conf=conf) sc_values[cls.__name__] = cls.sc @classmethod def tearDownClass(cls): print "....calling stop tearDownClas, the content of sc_values=", sc_values sc_values.clear() cls.sc.stop() class PySparkTestCase(unittest.TestCase): def setUp(self): self._old_sys_path = list(sys.path) conf = SparkConf().setMaster("local[2]") \ .setAppName(self.__class__.__name__) \ .set("spark.authenticate.secret", "111111") self.sc = SparkContext(conf=conf) def tearDown(self): self.sc.stop() sys.path = self._old_sys_path class TestResusedScA(ReusedPySparkTestCase): def testA_1(self): rdd = self.sc.parallelize([1,2,3]) self.assertEqual(rdd.collect(), [1,2,3]) sc_values['testA_1'] = self.sc def testA_2(self): sc_values['testA_2'] = self.sc self.assertEquals(self.sc, sc_values['testA_1']) class TestResusedScB(ReusedPySparkTestCase): def testB_1(self): sc_values['testB_1'] = self.sc def testB_2(self): sc_values['testB_2'] = self.sc def testB_3(self): sc_values['testB_3'] = self.sc self.assertEquals(self.sc, sc_values['testB_2']) if __name__ == '__main__': unittest.main() |
The output of the pyspark unit test
From the following output, we can see that all the test methods of TestResusedScA class share the same sparkContext, and all the test methods of TestResusedScB share the same sparkContext.
The output when call the tearDownClass for TestResusedScA:
|
1 |
......calling stop tearDownClas, the content of sc_values= {'testA_1': <pyspark.context.SparkContext object at 0x7f8ee1ed64d0>, 'TestResusedScA': <pyspark.context.SparkContext object at 0x7f8ee1ed64d0>, 'testA_2': <pyspark.context.SparkContext object at 0x7f8ee1ed64d0>} |
The output when call the tearDownClass for TestResusedScB:
|
1 |
.......calling stop tearDownClas, the content of sc_values= {'testB_1': <pyspark.context.SparkContext object at 0x7f8ee1e826d0>, 'testB_2': <pyspark.context.SparkContext object at 0x7f8ee1e826d0>, 'testB_3': <pyspark.context.SparkContext object at 0x7f8ee1e826d0>, 'TestResusedScB': <pyspark.context.SparkContext object at 0x7f8ee1e826d0>} |
Run the pyspark unit test
To run the above unit test for pyspark, we need to export the SPARK_HOME variable. Just run the following commands to start the pyspark unit test program.
|
1 2 3 |
export SPARK_HOME=/homes/.../spark python test.py |
In each of the test methods, as we can get the sparkContext reference by calling self.sc, we can conduct more complicated test using Spark RDD, and call self.assert* method to test our pyspark program.
A simple pyspark unit test example
In the following example, we develop a pyspark program to count the frequency of words in a set of sentences. Then we build a testClass to test the program.
testbase is the python module that contains the definition of the ReusedPySparkTestCase class.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
import testbase import unittest def wordCount(rdd): wcntRdd = rdd.flatMap(lambda line: line.split()).\ map(lambda word: (word, 1)).\ reduceByKey(lambda fa, fb: fa + fb) return wcntRdd class TestWordCount(testbase.ReusedPySparkTestCase): def test_word_count(self): rdd = self.sc.parallelize(["a b c d", "a c d e", "a d e f"]) res = wordCount(rdd) res = res.collectAsMap() expected = {"a":3, "b":1, "c":2, "d":3, "e":2, "f":1} self.assertEqual(res,expected) if __name__ == '__main__': if __name__ == '__main__': unittest.main() |
Reference:
The ReusedPySparkTestCase is defined in the following file of the spark github repository. You can refer to this file for more examples on how to do the pyspark unit test.
https://github.com/apache/spark/blob/master/python/pyspark/tests.py
-
Sven Hofstede
-
SidAli Ait
