Optional subfactories for Factory Boy

I often use a library called Factory Boy to create Django models in my tests. It works fine and I mostly like it but there's one thing that annoys me when it comes to nullable foreign keys on Django models. Let's set up some example models to explain:

from django.db import models

class Author(models.Model):
    name = models.CharField(max_length=100)
    email = models.EmailField(required=False)


class Article(models.Model):
    author = models.ForeignKey("Author", on_delete=models.CASCADE, blank=True, null=True)
    title = models.CharField(max_length=100)

Nothing extraordinary here, just two models linked by an optional foreign key. With Factory Boy (whose importable module is called factory for some reason), you create Factory classes for your models, and a foreign key relationship is typically represented using a SubFactory like so:

from factory.django import DjangoModelFactory
from factory import SubFactory

from .models import Article, Author

class AuthorFactory(DjangoModelFactory):
    class Meta:
        model = Author


class ArticleFactory(DjangoModelFactory):
    author = SubFactory(Author)

    class Meta:
        model = Article

That works great, and when you have those factories you can easily create two related instances in a single call to the factory:

>>> article = ArticleFactory.create(title="Test Article", author__name="Baptiste")
>>> article.author is None
False
>>> article.author.name
'Baptiste'

So far so expected. But a SubFactory is not "smart" enough to know that the foreign key is optional, and it will always created a related object, even if you don't ask for one:

>>> article = ArticleFactory.create(title="Test Article")
>>> article.author is None
False
>>> article.author.name
''

I'm not a big fan of this behavior, because I find it creates too many objects and often makes my tests more verbose than I'd like. Looking for a solution, I soon found a closed issue on github where the library author suggests using something called a Trait. But to be honest I never fully understood what traits were, and I always thought there had to be a simpler way.

After some hours of reading the code of the library and trying out several approaches, I ended up with something which I think works pretty well:

from factory.django import DjangoModelFactory
from factory import SubFactory

from .models import Article, Author


class OptionalSubFactory(SubFactory):
    def __init__(self, factory, **defaults):
        super().__init__(factory)
        self.defaults = defaults

    def evaluate(self, instance, step, extra):
        if not extra:
            return None
        kwargs = {**self.defaults, **extra}
        return super().evaluate(instance, step, kwargs)


class AuthorFactory(DjangoModelFactory):
    class Meta:
        model = Author


class ArticleFactory(DjangoModelFactory):
    author = OptionalSubFactory(Author)

    class Meta:
        model = Article

(In case you wanted to reuse this by any chance, the code in the above block is public domain.)

With this new OptionalSubFactory, related objects are created when you specify any of their arguments (just like with Factory):

>>> article = ArticleFactory.create(title="Test Article", author__name="Baptiste")
>>> article.author is None
False
>>> article.author.name
'Baptiste'

But if you leave out its arguments, then no object is created and the foreign key is set to None:

>>> article = ArticleFactory.create(title="Test Article")
>>> article.author is None
True
>>> article.author.name
Traceback (most recent call last):
  File "<console>", line 1, in <module>
AttributeError: 'NoneType' object has no attribute 'name'

Et voilà! The one downside I found and which I couldn't find a fix for is that you have to pass an explicit argument in order to create a related object. In some cases it would be nice to be able to write something like ArticleFactory.create(profile=True) and have the profile subfactory be called with all its default arguments. For now, if I want that behavior I have to do something like:

>>> ArticleFactory.create(profile=ProfileFactory.create())

That's a bit annoying, but it's a tradeoff I can live with.