From bf278fa3c7ca2b53d16f51c89fe79c60be89599d Mon Sep 17 00:00:00 2001
From: Haoyuan Li <haoyuan@cs.berkeley.edu>
Date: Sat, 22 Mar 2014 03:30:04 -0700
Subject: [PATCH] fix python tests

---
 python/pyspark/context.py      | 2 +-
 python/pyspark/rdd.py          | 3 ++-
 python/pyspark/storagelevel.py | 4 ++--
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/python/pyspark/context.py b/python/pyspark/context.py
index bf2454fd7e38e..efd942e4a23bf 100644
--- a/python/pyspark/context.py
+++ b/python/pyspark/context.py
@@ -383,7 +383,7 @@ def _getJavaStorageLevel(self, storageLevel):
             raise Exception("storageLevel must be of type pyspark.StorageLevel")
 
         newStorageLevel = self._jvm.org.apache.spark.storage.StorageLevel
-        return newStorageLevel(storageLevel.useDisk, storageLevel.useMemory,
+        return newStorageLevel(storageLevel.useDisk, storageLevel.useMemory, storageLevel.userTachyon,
             storageLevel.deserialized, storageLevel.replication)
 
     def setJobGroup(self, groupId, description):
diff --git a/python/pyspark/rdd.py b/python/pyspark/rdd.py
index 6d549b40e5698..e8aae4ecb665e 100644
--- a/python/pyspark/rdd.py
+++ b/python/pyspark/rdd.py
@@ -1178,11 +1178,12 @@ def getStorageLevel(self):
         Get the RDD's current storage level.
         >>> rdd1 = sc.parallelize([1,2])
         >>> rdd1.getStorageLevel()
-        StorageLevel(False, False, False, 1)
+        StorageLevel(False, False, False, False, 1)
         """
         java_storage_level = self._jrdd.getStorageLevel()
         storage_level = StorageLevel(java_storage_level.useDisk(),
                                      java_storage_level.useMemory(),
+                                     java_storage_level.useTachyon(),
                                      java_storage_level.deserialized(),
                                      java_storage_level.replication())
         return storage_level
diff --git a/python/pyspark/storagelevel.py b/python/pyspark/storagelevel.py
index 78af4d01a98b1..1f5e877f9e589 100644
--- a/python/pyspark/storagelevel.py
+++ b/python/pyspark/storagelevel.py
@@ -33,8 +33,8 @@ def __init__(self, useDisk, useMemory, useTachyon, deserialized, replication = 1
         self.replication = replication
 
     def __repr__(self):
-        return "StorageLevel(%s, %s, %s, %s)" % (
-            self.useDisk, self.useMemory, self.deserialized, self.replication)
+        return "StorageLevel(%s, %s, %s, %s, %s)" % (
+            self.useDisk, self.useMemory, self.useTachyon, self.deserialized, self.replication)
 
 StorageLevel.DISK_ONLY = StorageLevel(True, False, False, False)
 StorageLevel.DISK_ONLY_2 = StorageLevel(True, False, False, False, 2)