Add Factory.reset_sequence.

4 files changed, 139 insertions, 0 deletions

diff --git a/docs/changelog.rst b/docs/changelog.rst
index e820d8c..da3c4dc 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -13,6 +13,8 @@ ChangeLog
       fields defined in wrapping factories.
     - Move :class:`~factory.django.DjangoModelFactory` and :class:`~factory.mogo.MogoFactory`
       to their own modules (:mod:`factory.django` and :mod:`factory.mogo`)
+    - Add the :meth:`~factory.Factory.reset_sequence` classmethod to :class:`~factory.Factory`
+      to ease resetting the sequence counter for a given factory.
 
 *Deprecation:*
 
diff --git a/docs/reference.rst b/docs/reference.rst
index 377feb1..a2d6c9a 100644
--- a/docs/reference.rst
+++ b/docs/reference.rst
@@ -234,6 +234,48 @@ The :class:`Factory` class
         values, for instance.
 
 
+    **Advanced functions:**
+
+
+    .. classmethod:: reset_sequence(cls, value=None, force=False)
+
+        :arg int value: The value to reset the sequence to
+        :arg bool force: Whether to force-reset the sequence
+
+        Allows to reset the sequence counter for a :class:`~factory.Factory`.
+        The new value can be passed in as the ``value`` argument:
+
+        .. code-block:: pycon
+
+            >>> SomeFactory.reset_sequence(4)
+            >>> SomeFactory._next_sequence
+            4
+
+        Since subclasses of a non-:attr:`abstract <factory.Factory.ABSTRACT_FACTORY>`
+        :class:`~factory.Factory` share the same sequence counter, special care needs
+        to be taken when resetting the counter of such a subclass.
+
+        By default, :meth:`reset_sequence` will raise a :exc:`ValueError` when
+        called on a subclassed :class:`~factory.Factory` subclass. This can be
+        avoided by passing in the ``force=True`` flag:
+
+        .. code-block:: pycon
+
+            >>> InheritedFactory.reset_sequence()
+            Traceback (most recent call last):
+              File "factory_boy/tests/test_base.py", line 179, in test_reset_sequence_subclass_parent
+                SubTestObjectFactory.reset_sequence()
+              File "factory_boy/factory/base.py", line 250, in reset_sequence
+                "Cannot reset the sequence of a factory subclass. "
+            ValueError: Cannot reset the sequence of a factory subclass. Please call reset_sequence() on the root factory, or call reset_sequence(forward=True).
+
+            >>> InheritedFactory.reset_sequence(force=True)
+            >>>
+
+        This is equivalent to calling :meth:`reset_sequence` on the base
+        factory in the chain.
+
+
 .. _strategies:
 
 Strategies
diff --git a/factory/base.py b/factory/base.py
index 23fdac7..ff5404f 100644
--- a/factory/base.py
+++ b/factory/base.py
@@ -240,6 +240,21 @@ class BaseFactory(object):
     FACTORY_HIDDEN_ARGS = ()
 
     @classmethod
+    def reset_sequence(cls, value=None, force=False):
+        """Reset the sequence counter."""
+        if cls._base_factory:
+            if force:
+                cls._base_factory.reset_sequence(value=value)
+            else:
+                raise ValueError(
+                    "Cannot reset the sequence of a factory subclass. "
+                    "Please call reset_sequence() on the root factory, "
+                    "or call reset_sequence(forward=True)."
+                )
+        else:
+            cls._next_sequence = value
+
+    @classmethod
     def _setup_next_sequence(cls):
         """Set up an initial sequence value for Sequence attributes.
 
diff --git a/tests/test_base.py b/tests/test_base.py
index 73e59fa..4978d10 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -34,6 +34,7 @@ class TestObject(object):
         self.three = three
         self.four = four
 
+
 class FakeDjangoModel(object):
     @classmethod
     def create(cls, **kwargs):
@@ -46,6 +47,7 @@ class FakeDjangoModel(object):
             setattr(self, name, value)
             self.id = None
 
+
 class FakeModelFactory(base.Factory):
     ABSTRACT_FACTORY = True
 
@@ -116,6 +118,84 @@ class FactoryTestCase(unittest.TestCase):
         self.assertEqual(4, len(ones))
 
 
+class FactorySequenceTestCase(unittest.TestCase):
+    def setUp(self):
+        super(FactorySequenceTestCase, self).setUp()
+
+        class TestObjectFactory(base.Factory):
+            FACTORY_FOR = TestObject
+            one = declarations.Sequence(lambda n: n)
+
+        self.TestObjectFactory = TestObjectFactory
+
+    def test_reset_sequence(self):
+        o1 = self.TestObjectFactory()
+        self.assertEqual(0, o1.one)
+
+        o2 = self.TestObjectFactory()
+        self.assertEqual(1, o2.one)
+
+        self.TestObjectFactory.reset_sequence()
+        o3 = self.TestObjectFactory()
+        self.assertEqual(0, o3.one)
+
+    def test_reset_sequence_with_value(self):
+        o1 = self.TestObjectFactory()
+        self.assertEqual(0, o1.one)
+
+        o2 = self.TestObjectFactory()
+        self.assertEqual(1, o2.one)
+
+        self.TestObjectFactory.reset_sequence(42)
+        o3 = self.TestObjectFactory()
+        self.assertEqual(42, o3.one)
+
+    def test_reset_sequence_subclass_fails(self):
+        """Tests that the sequence of a 'slave' factory cannot be reseted."""
+        class SubTestObjectFactory(self.TestObjectFactory):
+            pass
+
+        self.assertRaises(ValueError, SubTestObjectFactory.reset_sequence)
+
+    def test_reset_sequence_subclass_force(self):
+        """Tests that reset_sequence(force=True) works."""
+        class SubTestObjectFactory(self.TestObjectFactory):
+            pass
+
+        o1 = SubTestObjectFactory()
+        self.assertEqual(0, o1.one)
+
+        o2 = SubTestObjectFactory()
+        self.assertEqual(1, o2.one)
+
+        SubTestObjectFactory.reset_sequence(force=True)
+        o3 = SubTestObjectFactory()
+        self.assertEqual(0, o3.one)
+
+        # The master sequence counter has been reset
+        o4 = self.TestObjectFactory()
+        self.assertEqual(1, o4.one)
+
+    def test_reset_sequence_subclass_parent(self):
+        """Tests that the sequence of a 'slave' factory cannot be reseted."""
+        class SubTestObjectFactory(self.TestObjectFactory):
+            pass
+
+        o1 = SubTestObjectFactory()
+        self.assertEqual(0, o1.one)
+
+        o2 = SubTestObjectFactory()
+        self.assertEqual(1, o2.one)
+
+        self.TestObjectFactory.reset_sequence()
+        o3 = SubTestObjectFactory()
+        self.assertEqual(0, o3.one)
+
+        o4 = self.TestObjectFactory()
+        self.assertEqual(1, o4.one)
+
+
+
 class FactoryDefaultStrategyTestCase(unittest.TestCase):
     def setUp(self):
         self.default_strategy = base.Factory.FACTORY_STRATEGY