Fix gc3libs.utils.YieldAtNext and its tests.

gc3pie · Jul 8, 2013 · 4129be5 · 4129be5
1 parent 841b4b6
commit 4129be5
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 19 deletions.
diff --git a/gc3pie/gc3libs/tests/test_utils.py b/gc3pie/gc3libs/tests/test_utils.py
@@ -25,6 +25,7 @@
 # stdlib imports
 import os
 import sys
+import unittest
 
 # 3rd party imports
 from nose.tools import assert_true, assert_false, assert_equal, raises
@@ -38,27 +39,44 @@
 
 ## test definitions
 
-def test_YieldAtNext():
+class TestYieldAtNext(object):
 
-    # create a generator that performs all steps in sequence
-    def yield_send_throw_seq():
-        # yield a value and get one back from `send`
-        val = (yield 0)
-        # yield that back and get an exception
-        try:
+    def test_YieldAtNext_yield(self):
+        def generator_yield():
+            yield 0
+        g = gc3libs.utils.YieldAtNext(generator_yield())
+        assert_equal(g.next(), 0)
+
+    def test_YieldAtNext_send(self):
+        def generator_yield_send():
+            val = (yield 1)
             yield val
-        except RuntimeError:
-            yield 2
+        g = gc3libs.utils.YieldAtNext(generator_yield_send())
+        assert_equal(g.next(), 1)
+        result = g.send('a sent value')
+        assert_equal(result, None)
+        assert_equal(g.next(), 'a sent value')
+
+    def test_YieldAtNext_throw(self):
+        def generator_yield_throw():
+            try:
+                val = (yield 2)
+            except RuntimeError:
+                yield 'exception caught'
+        g = gc3libs.utils.YieldAtNext(generator_yield_throw())
+        assert_equal(g.next(), 2)
+        result = g.throw(RuntimeError)
+        assert_equal(result, None)
+        assert_equal(g.next(), 'exception caught')
 
-    seq = gc3libs.utils.YieldAtNext(yield_send_throw_seq)
-    for i, val in enumerate(seq):
-        assert_equal(i, val)
-        if i == 0:
-            seq.send(1)
-        if i == 1:
-            seq.throw(RuntimeError)
-    assert_equal(i, 2)
-    assert_equal(val, 2)
+    @raises(StopIteration)
+    def test_YieldAtNext_StopIteration(self):
+        def generator_yield():
+            yield 3
+        g = gc3libs.utils.YieldAtNext(generator_yield())
+        assert_equal(g.next(), 3)
+        # raises `StopIteration`
+        g.next()
 
 
 ## main: run tests

diff --git a/gc3pie/gc3libs/utils.py b/gc3pie/gc3libs/utils.py
@@ -1483,6 +1483,9 @@ def __init__(self, generator):
         self._has_saved = False
         self._stop_iteration = False
 
+    def __iter__(self):
+        return self
+
     def next(self):
         if self._stop_iteration:
             raise StopIteration
@@ -1496,7 +1499,7 @@ def next(self):
 
     def send(self, value):
         try:
-            self._saved = self._generator.next()
+            self._saved = self._generator.send(value)
             self._has_saved = True
         except StopIteration:
             self._stop_iteration = True