diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index 0e5c1a3b3c2ad..a8d876d0fa3b3 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -487,9 +487,8 @@ def test_get_or_create(self):
         inputd = tempfile.mkdtemp()
         outputd = tempfile.mkdtemp() + "/"
 
-        def updater(it):
-            for k, vs, s in it:
-                yield (k, sum(vs, s or 0))
+        def updater(vs, s):
+            return sum(vs, s or 0)
 
         def setup():
             conf = SparkConf().set("spark.default.parallelism", 1)