[pypy-commit] pypy iterator-in-rpython: more tests, just in case

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: iterator-in-rpython
Changeset: r56008:60f42323dd92
Date: 2012-07-08 21:00 +0200
http://bitbucket.org/pypy/pypy/changeset/60f42323dd92/

Log:	more tests, just in case

diff --git a/pypy/rpython/test/test_rclass.py b/pypy/rpython/test/test_rclass.py
--- a/pypy/rpython/test/test_rclass.py
+++ b/pypy/rpython/test/test_rclass.py
@@ -1166,6 +1166,39 @@
 
         assert self.interpret(f, []) == f()
 
+    def test_iter_2_kinds(self):
+        class BaseIterable(object):
+            def __init__(self):
+                self.counter = 0
+            
+            def __iter__(self):
+                return self
+
+            def next(self):
+                if self.counter >= 5:
+                    raise StopIteration
+                self.counter += self.step
+                return self.counter - 1
+        
+        class Iterable(BaseIterable):
+            step = 1
+
+        class OtherIter(BaseIterable):
+            step = 2
+
+        def f(k):
+            if k:
+                i = Iterable()
+            else:
+                i = OtherIter()
+            s = 0
+            for elem in i:
+                s += elem
+            return s
+
+        assert self.interpret(f, [True]) == f(True)
+        assert self.interpret(f, [False]) == f(False)
+
 
 class TestOOtype(BaseTestRclass, OORtypeMixin):