Repository navigation
Expand file tree
/
Copy pathtestcase.py
More file actions
108 lines (88 loc) · 4.02 KB
/
Copy pathtestcase.py
File metadata and controls
108 lines (88 loc) · 4.02 KB
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
from __future__ import absolute_import, print_function
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Provides a common test case base for Python Spark tests"""
from .utils import add_pyspark_path, quiet_py4j
import unittest2
from pyspark.context import SparkContext
import os
class SparkTestingBaseTestCase(unittest2.TestCase):
"""Basic common test case for Spark. Provides a Spark context as sc.
For non local mode testing you can either override sparkMaster
or set the enviroment property SPARK_MASTER for non-local mode testing."""
@classmethod
def getMaster(cls):
return os.getenv('SPARK_MASTER', "local[4]")
def setUp(self):
"""Setup a basic Spark context for testing"""
self.sc = SparkContext(self.getMaster())
quiet_py4j()
def tearDown(self):
"""
Tear down the basic panda spark test case. This stops the running
context and does a hack to prevent Akka rebinding on the same port.
"""
self.sc.stop()
# To avoid Akka rebinding to the same port, since it doesn't unbind
# immediately on shutdown
self.sc._jvm.System.clearProperty("spark.driver.port")
def assertRDDEquals(self, expected, result):
return self.compareRDD(expected, result) == []
def compareRDD(self, expected, result):
expectedKeyed = expected.map(lambda x: (x, 1))\
.reduceByKey(lambda x, y: x + y)
resultKeyed = result.map(lambda x: (x, 1))\
.reduceByKey(lambda x, y: x + y)
return expectedKeyed.cogroup(resultKeyed)\
.map(lambda x: tuple(map(list, x[1])))\
.filter(lambda x: x[0] != x[1]).take(1)
def assertRDDEqualsWithOrder(self, expected, result):
return self.compareRDDWithOrder(expected, result) == []
def compareRDDWithOrder(self, expected, result):
def indexRDD(rdd):
return rdd.zipWithIndex().map(lambda x: (x[1], x[0]))
indexExpected = indexRDD(expected)
indexResult = indexRDD(result)
return indexExpected.cogroup(indexResult)\
.map(lambda x: tuple(map(list, x[1])))\
.filter(lambda x: x[0] != x[1]).take(1)
class SparkTestingBaseReuse(unittest2.TestCase):
"""Basic common test case for Spark. Provides a Spark context as sc.
For non local mode testing you can either override sparkMaster
or set the enviroment property SPARK_MASTER for non-local mode testing."""
@classmethod
def getMaster(cls):
return os.getenv('SPARK_MASTER', "local[4]")
@classmethod
def setUpClass(cls):
"""Setup a basic Spark context for testing"""
class_name = cls.__name__
cls.sc = SparkContext(cls.getMaster(), appName=class_name)
quiet_py4j()
@classmethod
def tearDownClass(cls):
"""
Tear down the basic panda spark test case. This stops the running
context and does a hack to prevent Akka rebinding on the same port.
"""
print("stopping class")
cls.sc.stop()
# To avoid Akka rebinding to the same port, since it doesn't unbind
# immediately on shutdown
cls.sc._jvm.System.clearProperty("spark.driver.port")
if __name__ == "__main__":
unittest2.main()