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

1 files changed, 106 insertions, 0 deletions

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))
+