diff options
Diffstat (limited to 'python/pyspark/mllib/random.py')
| -rw-r--r-- | python/pyspark/mllib/random.py | 10 |
1 files changed, 7 insertions, 3 deletions
diff --git a/python/pyspark/mllib/random.py b/python/pyspark/mllib/random.py index 6a3c643b66..61213ddf62 100644 --- a/python/pyspark/mllib/random.py +++ b/python/pyspark/mllib/random.py @@ -409,13 +409,17 @@ class RandomRDDs(object): def _test(): import doctest - from pyspark.context import SparkContext + from pyspark.sql import SparkSession globs = globals().copy() # The small batch size here ensures that we see multiple batches, # even in these small test examples: - globs['sc'] = SparkContext('local[2]', 'PythonTest', batchSize=2) + spark = SparkSession.builder\ + .master("local[2]")\ + .appName("mllib.random tests")\ + .getOrCreate() + globs['sc'] = spark.sparkContext (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) - globs['sc'].stop() + spark.stop() if failure_count: exit(-1) |