diff options
Diffstat (limited to 'python/pyspark/mllib/fpm.py')
| -rw-r--r-- | python/pyspark/mllib/fpm.py | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/python/pyspark/mllib/fpm.py b/python/pyspark/mllib/fpm.py index f339e50891..ab4066f7d6 100644 --- a/python/pyspark/mllib/fpm.py +++ b/python/pyspark/mllib/fpm.py @@ -183,16 +183,21 @@ class PrefixSpan(object): def _test(): import doctest + from pyspark.sql import SparkSession import pyspark.mllib.fpm globs = pyspark.mllib.fpm.__dict__.copy() - globs['sc'] = SparkContext('local[4]', 'PythonTest') + spark = SparkSession.builder\ + .master("local[4]")\ + .appName("mllib.fpm tests")\ + .getOrCreate() + globs['sc'] = spark.sparkContext import tempfile temp_path = tempfile.mkdtemp() globs['temp_path'] = temp_path try: (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) - globs['sc'].stop() + spark.stop() finally: from shutil import rmtree try: |