diff --git a/factory/base.py b/factory/base.py index 36b2359a..77056a1b 100644 --- a/factory/base.py +++ b/factory/base.py @@ -664,8 +664,23 @@ def create(cls, **kwargs): raise errors.UnsupportedStrategy() +class DictFactoryOptions(FactoryOptions): + + def _get_counter_reference(self): + """Identify which factory should be used for a shared counter.""" + + # When _setup_next_sequence is overridden in the factory, it should get a new counter + if '_setup_next_sequence' in self.factory.__dict__: + return self + else: + return super()._get_counter_reference() + + class BaseDictFactory(Factory): """Factory for dictionary-like classes.""" + + _options_class = DictFactoryOptions + class Meta: abstract = True diff --git a/tests/test_base.py b/tests/test_base.py index d3b32570..8a8ad442 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -445,6 +445,40 @@ class Meta: self.assertEqual(enums.CREATE_STRATEGY, TestModelFactory._meta.strategy) +class DictFactoryTestCase(unittest.TestCase): + def setUp(self): + super().setUp() + pet_sequence_offset = 10 + + class Pet(base.DictFactory): + @classmethod + def _setup_next_sequence(cls) -> int: + return pet_sequence_offset + + pet_id = declarations.Sequence(lambda n: n) + + class Cat(Pet): + pass + + class Dog(Pet): + pass + + self.Pet = Pet + self.Cat = Cat + self.Dog = Dog + self.pet_sequence_offset = pet_sequence_offset + + def test_override_setup_next_sequence(self): + result = self.Pet() + self.assertEqual(result['pet_id'], self.pet_sequence_offset) + + def test_descendants_share_counter(self): + cat_result = self.Cat() + dog_result = self.Dog() + results = [cat_result['pet_id'], dog_result['pet_id']] + self.assertEqual(results, [self.pet_sequence_offset, self.pet_sequence_offset + 1]) + + class FactoryCreationTestCase(unittest.TestCase): def test_factory_for(self): class TestFactory(base.Factory):