diff --git a/autofixture/base.py b/autofixture/base.py index 3c482f1..a02b5de 100644 --- a/autofixture/base.py +++ b/autofixture/base.py @@ -1,10 +1,11 @@ # -*- coding: utf-8 -*- import inspect import warnings +import copy from django.db import models from django.db.models import fields from django.db.models.fields import related -from django.contrib.contenttypes.generic import GenericRelation +from django.contrib.contenttypes.generic import GenericRelation, GenericForeignKey from django.utils.datastructures import SortedDict from django.utils.six import with_metaclass import autofixture @@ -93,6 +94,7 @@ class IGNORE_FIELD(object): generate_fk = False follow_m2m = {'ALL': (1,5)} generate_m2m = False + generate_genericfk = False none_p = 0.2 tries = 1000 @@ -110,6 +112,7 @@ class IGNORE_FIELD(object): (fields.IPAddressField, generators.IPAddressGenerator), (fields.TextField, generators.LoremGenerator), (fields.TimeField, generators.TimeGenerator), + (GenericForeignKey, generators.GenericFKSelector) )) field_values = Values() @@ -121,7 +124,7 @@ class IGNORE_FIELD(object): def __init__(self, model, field_values=None, none_p=None, overwrite_defaults=None, constraints=None, follow_fk=None, generate_fk=None, - follow_m2m=None, generate_m2m=None): + follow_m2m=None, generate_m2m=None, generate_genericfk=None): ''' Parameters: ``model``: A model class which is used to create the test data. @@ -201,7 +204,9 @@ def __init__(self, model, self.generate_m2m = generate_m2m if not isinstance(self.generate_m2m, Link): self.generate_m2m = Link(self.generate_m2m) - + if generate_genericfk is not None: + self.generate_genericfk = generate_genericfk + for constraint in self.default_constraints: self.add_constraint(constraint) @@ -230,6 +235,18 @@ def add_constraint(self, constraint): ''' self.constraints.append(constraint) + def _normalize_genericfk_field(self, field): + """ + Add some attributes to the GenericFK field so that it behaves more + like "regular" fields and the usual checks don't fail. + """ + field_copy = copy.copy(field) + field_copy.default = fields.NOT_PROVIDED + fk_field_name = field.fk_field + field_copy.null = self.model._meta.get_field_by_name(fk_field_name)[0].null + field_copy.choices = [] + return field_copy + def get_generator(self, field): ''' Return a value generator based on the field instance that is passed to @@ -237,6 +254,9 @@ def get_generator(self, field): specified field will be ignored (e.g. if no matching generator was found). ''' + if isinstance(field, GenericForeignKey): + field = self._normalize_genericfk_field(field) + if isinstance(field, fields.AutoField): return None if isinstance(field, related.OneToOneField) and field.primary_key: @@ -247,7 +267,6 @@ def get_generator(self, field): field.name not in self.field_values): return None kwargs = {} - if field.name in self.field_values: value = self.field_values[field.name] if isinstance(value, generators.Generator): @@ -356,6 +375,8 @@ def get_generator(self, field): min_value=-field.MAX_BIGINT - 1, max_value=field.MAX_BIGINT, **kwargs) + if isinstance(field, GenericForeignKey): + return generators.GenericFKSelector(generate_genericfk=self.generate_genericfk) for field_class, generator in self.field_to_generator.items(): if isinstance(field, field_class): return generator(**kwargs) @@ -461,7 +482,16 @@ def create_one(self, commit=True): ''' tries = self.tries instance = self.model() - process = instance._meta.fields + process = copy.copy(instance._meta.fields) + + #remove genericfk field components, add virtualfield instead + generic_fields = instance._meta.virtual_fields + for field in generic_fields: + if isinstance(field, GenericForeignKey): + process.append(field) + ct_field = instance._meta.get_field(field.ct_field) + fk_field = instance._meta.get_field(field.fk_field) + process = [f for f in process if not f in [ct_field, fk_field]] while process and tries > 0: for field in process: self.process_field(instance, field) diff --git a/autofixture/generators.py b/autofixture/generators.py index e64d515..b24c1e6 100644 --- a/autofixture/generators.py +++ b/autofixture/generators.py @@ -36,6 +36,10 @@ def relpath(path, start=os.curdir): return os.path.join(*rel_list) +class GeneratorError(Exception): + pass + + class Generator(object): coerce_type = staticmethod(lambda x: x) empty_value = None @@ -582,3 +586,48 @@ def weighted_choice(self, choices): def generate(self): return self.weighted_choice(self.choices).generate() + +class GenericFKSelector(Generator): + """ + Should return an instance of some object. + """ + def __init__(self, generate_genericfk=False, limit_ct_to=None, + limit_ids_to=None): + self.generate_genericfk = generate_genericfk + self.limit_ct_to = limit_ct_to or {} + self.limit_ids_to = limit_ids_to + + def get_ct(self): + """ + Get content type object + """ + from django.contrib.contenttypes.models import ContentType + queryset = ContentType.objects.filter(**self.limit_ct_to) + if not queryset: + raise GeneratorError( + "Found no contenttypes for filter params %s" %self.limit_ct_to) + #get any old ct, we'll generate objects later + if self.generate_genericfk: + return InstanceSelector(queryset=queryset).generate() + else: # find a contenttype with some existing objects + for ct in queryset.order_by("?"): + if ct.get_all_objects_for_this_type().count() > 0: + return ct + raise GeneratorError( + """Found no contenttypes for filter params %s + that have already existing objects""" %self.limit_ct_to ) + + def get_object(self, content_type): + # if option 'generate_genericfk' + queryset = content_type.get_all_objects_for_this_type() + if self.generate_genericfk: + return InstanceGenerator(autofixture=AutoFixture( + content_type.model_class()), + limit_choices_to=self.limit_ids_to).generate() + else: + return InstanceSelector(queryset=queryset, + limit_choices_to=self.limit_ids_to).generate() + + def generate(self): + ct = self.get_ct() + return self.get_object(ct)