move unit test to tests.py

python/pyspark/mllib/fpm.py

@@ -36,16 +36,13 @@ class FPGrowthModel(JavaModelWrapper, JavaSaveable, JavaLoader):
 
     >>> data = [["a", "b", "c"], ["a", "b", "d", "e"], ["a", "c", "e"], ["a", "c", "f"]]
     >>> rdd = sc.parallelize(data, 2)
-    >>> model1 = FPGrowth.train(rdd, 0.6, 2)
-    >>> model2 = FPGrowth.train(rdd, 0.6)
-    >>> sorted(model1.freqItemsets().collect())
-    [FreqItemset(items=[u'a'], freq=4), FreqItemset(items=[u'c'], freq=3), ...
-    >>> sorted(model2.freqItemsets().collect())
+    >>> model = FPGrowth.train(rdd, 0.6, 2)
+    >>> sorted(model.freqItemsets().collect())
     [FreqItemset(items=[u'a'], freq=4), FreqItemset(items=[u'c'], freq=3), ...
     >>> model_path = temp_path + "/fpm"
-    >>> model1.save(sc, model_path)
+    >>> model.save(sc, model_path)
     >>> sameModel = FPGrowthModel.load(sc, model_path)
-    >>> sorted(model1.freqItemsets().collect()) == sorted(sameModel.freqItemsets().collect())
+    >>> sorted(model.freqItemsets().collect()) == sorted(sameModel.freqItemsets().collect())
     True
 
     .. versionadded:: 1.4.0

python/pyspark/mllib/tests.py

@@ -57,6 +57,7 @@
     DenseMatrix, SparseMatrix, Vectors, Matrices, MatrixUDT
 from pyspark.mllib.linalg.distributed import RowMatrix
 from pyspark.mllib.classification import StreamingLogisticRegressionWithSGD
+from pyspark.mllib.fpm import FPGrowth
 from pyspark.mllib.recommendation import Rating
 from pyspark.mllib.regression import LabeledPoint, StreamingLinearRegressionWithSGD
 from pyspark.mllib.random import RandomRDDs
@@ -1762,6 +1763,17 @@ def test_pca(self):
                 self.assertEqualUpToSign(pcs.toArray()[:, k - 1], expected_pcs[:, k - 1])
 
 
+class FPGrowthTest(MLlibTestCase):
+
+    def test_fpgrowth(self):
+        data = [["a", "b", "c"], ["a", "b", "d", "e"], ["a", "c", "e"], ["a", "c", "f"]]
+        rdd = self.sc.parallelize(data, 2)
+        model1 = FPGrowth.train(rdd, 0.6, 2)
+        # use default data partition number when numPartitions is not specified
+        model2 = FPGrowth.train(rdd, 0.6)
+        self.assertEqual(sorted(model1.freqItemsets().collect()),
+                         sorted(model2.freqItemsets().collect()))
+
 if __name__ == "__main__":
     from pyspark.mllib.tests import *
     if not _have_scipy: