Add the Iterator and InfiniteIterator attribute kinds.

Signed-off-by: Raphaël Barrois <raphael.barrois@polytechnique.org>

1 files changed, 46 insertions, 0 deletions

diff --git a/tests/test_using.py b/tests/test_using.py
index 4e69212..41b34ea 100644
--- a/tests/test_using.py
+++ b/tests/test_using.py
@@ -760,5 +760,51 @@ class SubFactoryTestCase(unittest.TestCase):
         self.assertEqual(outer.side_a.inner_from_a.b, 4)
 
 
+class IteratorTestCase(unittest.TestCase):
+
+    def test_iterator(self):
+        class TestObjectFactory(factory.Factory):
+            one = factory.Iterator(xrange(10, 30))
+
+        objs = TestObjectFactory.build_batch(20)
+
+        for i, obj in enumerate(objs):
+            self.assertEqual(i + 10, obj.one)
+
+    def test_infinite_iterator(self):
+        class TestObjectFactory(factory.Factory):
+            one = factory.InfiniteIterator(xrange(5))
+
+        objs = TestObjectFactory.build_batch(20)
+
+        for i, obj in enumerate(objs):
+            self.assertEqual(i % 5, obj.one)
+
+    def test_iterator_decorator(self):
+        class TestObjectFactory(factory.Factory):
+            @factory.iterator
+            def one():
+                for i in xrange(10, 50):
+                    yield i
+
+        objs = TestObjectFactory.build_batch(20)
+
+        for i, obj in enumerate(objs):
+            self.assertEqual(i + 10, obj.one)
+
+    def test_infinite_iterator_decorator(self):
+        class TestObjectFactory(factory.Factory):
+            @factory.infinite_iterator
+            def one():
+                for i in xrange(5):
+                    yield i
+
+        objs = TestObjectFactory.build_batch(20)
+
+        for i, obj in enumerate(objs):
+            self.assertEqual(i % 5, obj.one)
+
+
+
 if __name__ == '__main__':
     unittest.main()