AutoFixture for GenericForeignKey

fixtures.py

from django.contrib.contenttypes.fields import GenericForeignKey
from django.contrib.contenttypes.models import ContentType

from autofixture import AutoFixture, generators


class GenericFKAutoFixture(AutoFixture):
    """
    AutoFixture class that generate instances with ::GenericForeignKey:: that
    make sense. This is, the ::object_id:: will match an acutal instance of the
    selected ::content_type::

    You need to subclass this AutoFixture and set the ::GENERIC_FK_MAP::.
    The keys shoud match the GenericForeignKeys in your model.
    The values are a dict where entries have the form <model>: <generator>.

    e.g.

    GENERIC_FK_MAP = {
        generic_foreign_key_field_1: {
            model_a: generator_a,
            model_b: generator_b,
            ...
        },
        generic_foreign_key_field_2: {
            model_c: generator_c,
            model_d: generator_d,
            ...
        },
        ...
    }
    """
    GENERIC_FK_MAP = {}

    def __init__(self, model, *args, **kwargs):
        super(GenericFKAutoFixture, self).__init__(model, *args, **kwargs)

        for field in model._meta.virtual_fields:
            if isinstance(field, GenericForeignKey):
                self._prepate_generic_fk_field(field)

    def _prepate_generic_fk_field(self, field):
        """
        Given a ::GenericForeignKey:: field without an ::field_value::
        specified set the value to an ::InstanceSelector:: generator, limiting
        the options to the models specified in the keys of ::GENERIC_FK_MAP::
        """
        field_value = self.field_values.get(field.ct_field, None)
        if field_value is None:
            choices = self.GENERIC_FK_MAP.get(field.name, {})
            choices = [m._meta.model_name for m in choices.keys()]
            field_value = generators.InstanceSelector(
                ContentType,
                limit_choices_to={'model__in': choices}
            )
            self.add_field_value(field.ct_field, field_value)

    def pre_process_instance(self, instance):
        """
        Given an instance, search the ::GenericForeignKey:: fields and set the
        ::object_id:: field to the pk of an instance of the type specified by
        the ::content_type:: field.
        """
        for field in self.model._meta.virtual_fields:
            if isinstance(field, GenericForeignKey):
                generic_model = getattr(instance, field.ct_field).model_class()
                choices = self.GENERIC_FK_MAP.get(field.name, {})
                selector = choices.get(generic_model, None)
                if selector is not None:
                    source = selector.generate()
                    value = source.pk if source is not None else None
                    setattr(instance, field.fk_field, value)

        return instance