test view.map on a generator

IPython/parallel/tests/test_lbview.py

@@ -1,203 +1,211 b''
 # -*- coding: utf-8 -*-
 """test LoadBalancedView objects
 
 Authors:
 
 * Min RK
 """
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
 import sys
 import time
 
 import zmq
 from nose import SkipTest
 from nose.plugins.attrib import attr
 
 from IPython import parallel  as pmod
 from IPython.parallel import error
 
 from IPython.parallel.tests import add_engines
 
 from .clienttest import ClusterTestCase, crash, wait, skip_without
 
 def setup():
     add_engines(3, total=True)
 
 class TestLoadBalancedView(ClusterTestCase):
 
     def setUp(self):
         ClusterTestCase.setUp(self)
         self.view = self.client.load_balanced_view()
 
     @attr('crash')
     def test_z_crash_task(self):
         """test graceful handling of engine death (balanced)"""
         # self.add_engines(1)
         ar = self.view.apply_async(crash)
         self.assertRaisesRemote(error.EngineError, ar.get, 10)
         eid = ar.engine_id
         tic = time.time()
         while eid in self.client.ids and time.time()-tic < 5:
             time.sleep(.01)
             self.client.spin()
         self.assertFalse(eid in self.client.ids, "Engine should have died")
 
     def test_map(self):
         def f(x):
             return x**2
         data = range(16)
         r = self.view.map_sync(f, data)
         self.assertEqual(r, map(f, data))
-
+    
+    def test_map_generator(self):
+        def f(x):
+            return x**2
+        
+        data = range(16)
+        r = self.view.map_sync(f, iter(data))
+        self.assertEqual(r, map(f, iter(data)))
+    
     def test_map_short_first(self):
         def f(x,y):
             if y is None:
                 return y
             if x is None:
                 return x
             return x*y
         data = range(10)
         data2 = range(4)
         
         r = self.view.map_sync(f, data, data2)
         self.assertEqual(r, map(f, data, data2))
 
     def test_map_short_last(self):
         def f(x,y):
             if y is None:
                 return y
             if x is None:
                 return x
             return x*y
         data = range(4)
         data2 = range(10)
         
         r = self.view.map_sync(f, data, data2)
         self.assertEqual(r, map(f, data, data2))
 
     def test_map_unordered(self):
         def f(x):
             return x**2
         def slow_f(x):
             import time
             time.sleep(0.05*x)
             return x**2
         data = range(16,0,-1)
         reference = map(f, data)
         
         amr = self.view.map_async(slow_f, data, ordered=False)
         self.assertTrue(isinstance(amr, pmod.AsyncMapResult))
         # check individual elements, retrieved as they come
         # list comprehension uses __iter__
         astheycame = [ r for r in amr ]
         # Ensure that at least one result came out of order:
         self.assertNotEqual(astheycame, reference, "should not have preserved order")
         self.assertEqual(sorted(astheycame, reverse=True), reference, "result corrupted")
 
     def test_map_ordered(self):
         def f(x):
             return x**2
         def slow_f(x):
             import time
             time.sleep(0.05*x)
             return x**2
         data = range(16,0,-1)
         reference = map(f, data)
         
         amr = self.view.map_async(slow_f, data)
         self.assertTrue(isinstance(amr, pmod.AsyncMapResult))
         # check individual elements, retrieved as they come
         # list(amr) uses __iter__
         astheycame = list(amr)
         # Ensure that results came in order
         self.assertEqual(astheycame, reference)
         self.assertEqual(amr.result, reference)
 
     def test_map_iterable(self):
         """test map on iterables (balanced)"""
         view = self.view
         # 101 is prime, so it won't be evenly distributed
         arr = range(101)
         # so that it will be an iterator, even in Python 3
         it = iter(arr)
         r = view.map_sync(lambda x:x, arr)
         self.assertEqual(r, list(arr))
 
     
     def test_abort(self):
         view = self.view
         ar = self.client[:].apply_async(time.sleep, .5)
         ar = self.client[:].apply_async(time.sleep, .5)
         time.sleep(0.2)
         ar2 = view.apply_async(lambda : 2)
         ar3 = view.apply_async(lambda : 3)
         view.abort(ar2)
         view.abort(ar3.msg_ids)
         self.assertRaises(error.TaskAborted, ar2.get)
         self.assertRaises(error.TaskAborted, ar3.get)
 
     def test_retries(self):
         view = self.view
         view.timeout = 1 # prevent hang if this doesn't behave
         def fail():
             assert False
         for r in range(len(self.client)-1):
             with view.temp_flags(retries=r):
                 self.assertRaisesRemote(AssertionError, view.apply_sync, fail)
 
         with view.temp_flags(retries=len(self.client), timeout=0.25):
             self.assertRaisesRemote(error.TaskTimeout, view.apply_sync, fail)
 
     def test_invalid_dependency(self):
         view = self.view
         with view.temp_flags(after='12345'):
             self.assertRaisesRemote(error.InvalidDependency, view.apply_sync, lambda : 1)
 
     def test_impossible_dependency(self):
         self.minimum_engines(2)
         view = self.client.load_balanced_view()
         ar1 = view.apply_async(lambda : 1)
         ar1.get()
         e1 = ar1.engine_id
         e2 = e1
         while e2 == e1:
             ar2 = view.apply_async(lambda : 1)
             ar2.get()
             e2 = ar2.engine_id
 
         with view.temp_flags(follow=[ar1, ar2]):
             self.assertRaisesRemote(error.ImpossibleDependency, view.apply_sync, lambda : 1)
 
 
     def test_follow(self):
         ar = self.view.apply_async(lambda : 1)
         ar.get()
         ars = []
         first_id = ar.engine_id
 
         self.view.follow = ar
         for i in range(5):
             ars.append(self.view.apply_async(lambda : 1))
         self.view.wait(ars)
         for ar in ars:
             self.assertEqual(ar.engine_id, first_id)
 
     def test_after(self):
         view = self.view
         ar = view.apply_async(time.sleep, 0.5)
         with view.temp_flags(after=ar):
             ar2 = view.apply_async(lambda : 1)
 
         ar.wait()
         ar2.wait()
         self.assertTrue(ar2.started >= ar.completed, "%s not >= %s"%(ar.started, ar.completed))