[SPARK-8511][pyspark] Modify a test to remove a saved model in `regre…

…ssion.py`
maropu · Jun 21, 2015 · 4a01c9e · 4a01c9e
1 parent a1e3649
commit 4a01c9e
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/python/pyspark/mllib/regression.py b/python/pyspark/mllib/regression.py
@@ -133,8 +133,9 @@ class LinearRegressionModel(LinearRegressionModelBase):
     True
     >>> abs(sameModel.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5
     True
+    >>> from shutil import rmtree
     >>> try:
-    ...    os.removedirs(path)
+    ...    rmtree(path)
     ... except:
     ...    pass
     >>> data = [
@@ -275,8 +276,9 @@ class LassoModel(LinearRegressionModelBase):
     True
     >>> abs(sameModel.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5
     True
+    >>> from shutil import rmtree
     >>> try:
-    ...    os.removedirs(path)
+    ...    rmtree(path)
     ... except:
     ...    pass
     >>> data = [
@@ -389,8 +391,9 @@ class RidgeRegressionModel(LinearRegressionModelBase):
     True
     >>> abs(sameModel.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5
     True
+    >>> from shutil import rmtree
     >>> try:
-    ...    os.removedirs(path)
+    ...    rmtree(path)
     ... except:
     ...    pass
     >>> data = [