Skip to content

Commit

Permalink
Respect default manager in DjangoModelFactory (Closes #192).
Browse files Browse the repository at this point in the history
The previous version tries to use ``cls._default_manager`` all the time,
which breaks with ``manager.using(db_name)``.
  • Loading branch information
rbarrois committed Mar 27, 2015
1 parent 03ca4ec commit bdc1b81
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 7 deletions.
3 changes: 3 additions & 0 deletions docs/changelog.rst
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@ ChangeLog
2.5.1 (master)
--------------

*Bugfix:*

- Respect custom managers in :class:`~factory.django.DjangoModelFactory` (see :issue:`192`)

.. _v2.5.0:

Expand Down
13 changes: 6 additions & 7 deletions factory/django.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@
logger = logging.getLogger('factory.generate')


DEFAULT_DB_ALIAS = 'default' # Same as django.db.DEFAULT_DB_ALIAS


def require_django():
"""Simple helper to ensure Django is available."""
Expand All @@ -56,7 +58,7 @@ class DjangoOptions(base.FactoryOptions):
def _build_default_options(self):
return super(DjangoOptions, self)._build_default_options() + [
base.OptionDefault('django_get_or_create', (), inherit=True),
base.OptionDefault('database', 'default', inherit=True),
base.OptionDefault('database', DEFAULT_DB_ALIAS, inherit=True),
]

def _get_counter_reference(self):
Expand Down Expand Up @@ -100,12 +102,9 @@ def _get_manager(cls, model_class):
if model_class is None:
raise base.AssociatedClassError("No model set on %s.%s.Meta"
% (cls.__module__, cls.__name__))
try:
manager = model_class._default_manager # pylint: disable=W0212
except AttributeError:
manager = model_class.objects

manager = manager.using(cls._meta.database)
manager = model_class.objects
if cls._meta.database != DEFAULT_DB_ALIAS:
manager = manager.using(cls._meta.database)
return manager

@classmethod
Expand Down
17 changes: 17 additions & 0 deletions tests/djapp/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,3 +87,20 @@ class WithImage(models.Model):

class WithSignals(models.Model):
foo = models.CharField(max_length=20)


class CustomQuerySet(models.QuerySet):
pass


class CustomManager(models.Manager):

def create(self, arg=None, **kwargs):
return super(CustomManager, self).create(**kwargs)


class WithCustomManager(models.Model):

foo = models.CharField(max_length=20)

objects = CustomManager.from_queryset(CustomQuerySet)()
12 changes: 12 additions & 0 deletions tests/test_django.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,13 @@ class Meta:
model = models.WithSignals


class WithCustomManagerFactory(factory.django.DjangoModelFactory):
class Meta:
model = models.WithCustomManager

foo = factory.Sequence(lambda n: "foo%d" % n)


@unittest.skipIf(django is None, "Django not installed.")
class ModelTests(django_test.TestCase):
def test_unset_model(self):
Expand Down Expand Up @@ -706,5 +713,10 @@ def generate(cls):

self.assertSignalsReactivated()

class DjangoCustomManagerTestCase(django_test.TestCase):

def test_extra_args(self):
model = WithCustomManagerFactory(arg='foo')

if __name__ == '__main__': # pragma: no cover
unittest.main()

0 comments on commit bdc1b81

Please sign in to comment.