Add Iterator.reset() (Closes #63).

2 files changed, 127 insertions, 0 deletions

diff --git a/tests/test_declarations.py b/tests/test_declarations.py
index cd38dd2..9d54c59 100644
--- a/tests/test_declarations.py
+++ b/tests/test_declarations.py
@@ -107,6 +107,27 @@ class IteratorTestCase(unittest.TestCase):
         self.assertEqual(2, it.evaluate(1, None, False))
         self.assertRaises(StopIteration, it.evaluate, 2, None, False)
 
+    def test_reset_cycle(self):
+        it = declarations.Iterator([1, 2])
+        self.assertEqual(1, it.evaluate(0, None, False))
+        self.assertEqual(2, it.evaluate(1, None, False))
+        self.assertEqual(1, it.evaluate(2, None, False))
+        self.assertEqual(2, it.evaluate(3, None, False))
+        self.assertEqual(1, it.evaluate(4, None, False))
+        it.reset()
+        self.assertEqual(1, it.evaluate(5, None, False))
+        self.assertEqual(2, it.evaluate(6, None, False))
+
+    def test_reset_no_cycling(self):
+        it = declarations.Iterator([1, 2], cycle=False)
+        self.assertEqual(1, it.evaluate(0, None, False))
+        self.assertEqual(2, it.evaluate(1, None, False))
+        self.assertRaises(StopIteration, it.evaluate, 2, None, False)
+        it.reset()
+        self.assertEqual(1, it.evaluate(0, None, False))
+        self.assertEqual(2, it.evaluate(1, None, False))
+        self.assertRaises(StopIteration, it.evaluate, 2, None, False)
+
     def test_getter(self):
         it = declarations.Iterator([(1, 2), (1, 3)], getter=lambda p: p[1])
         self.assertEqual(2, it.evaluate(0, None, False))
diff --git a/tests/test_utils.py b/tests/test_utils.py
index b353c9d..8c73935 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
+import itertools
 
 from factory import utils
 
@@ -230,3 +231,108 @@ class ImportObjectTestCase(unittest.TestCase):
     def test_invalid_module(self):
         self.assertRaises(ImportError, utils.import_object,
             'this-is-an-invalid-module', '__name__')
+
+
+class ResetableIteratorTestCase(unittest.TestCase):
+    def test_no_reset(self):
+        i = utils.ResetableIterator([1, 2, 3])
+        self.assertEqual([1, 2, 3], list(i))
+
+    def test_no_reset_new_iterator(self):
+        i = utils.ResetableIterator([1, 2, 3])
+        iterator = iter(i)
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+
+        iterator2 = iter(i)
+        self.assertEqual(3, next(iterator2))
+
+    def test_infinite(self):
+        i = utils.ResetableIterator(itertools.cycle([1, 2, 3]))
+        iterator = iter(i)
+        values = [next(iterator) for _i in range(10)]
+        self.assertEqual([1, 2, 3, 1, 2, 3, 1, 2, 3, 1], values)
+
+    def test_reset_simple(self):
+        i = utils.ResetableIterator([1, 2, 3])
+        iterator = iter(i)
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+
+        i.reset()
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+
+    def test_reset_at_begin(self):
+        i = utils.ResetableIterator([1, 2, 3])
+        iterator = iter(i)
+        i.reset()
+        i.reset()
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+
+    def test_reset_at_end(self):
+        i = utils.ResetableIterator([1, 2, 3])
+        iterator = iter(i)
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+
+        i.reset()
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+
+    def test_reset_after_end(self):
+        i = utils.ResetableIterator([1, 2, 3])
+        iterator = iter(i)
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+        self.assertRaises(StopIteration, next, iterator)
+
+        i.reset()
+        # Previous iter() has stopped
+        iterator = iter(i)
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+
+    def test_reset_twice(self):
+        i = utils.ResetableIterator([1, 2, 3, 4, 5])
+        iterator = iter(i)
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+
+        i.reset()
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+        self.assertEqual(4, next(iterator))
+
+        i.reset()
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+        self.assertEqual(4, next(iterator))
+
+    def test_reset_shorter(self):
+        i = utils.ResetableIterator([1, 2, 3, 4, 5])
+        iterator = iter(i)
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+        self.assertEqual(4, next(iterator))
+
+        i.reset()
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+
+        i.reset()
+        self.assertEqual(1, next(iterator))
+        self.assertEqual(2, next(iterator))
+        self.assertEqual(3, next(iterator))
+        self.assertEqual(4, next(iterator))
+