diff --git a/django_seed/seeder.py b/django_seed/seeder.py index 07d8e97..a2e487a 100644 --- a/django_seed/seeder.py +++ b/django_seed/seeder.py @@ -1,12 +1,13 @@ import random -from django_seed.guessers import NameGuesser, FieldTypeGuesser -from django_seed.exceptions import SeederException -from django.db.models.fields import AutoField + from django.db.models import ForeignKey, ManyToManyField, OneToOneField +from django.db.models.fields import AutoField +from django_seed.exceptions import SeederException +from django_seed.guessers import NameGuesser, FieldTypeGuesser -class ModelSeeder(object): +class ModelSeeder(object): def __init__(self, model): """ :param model: Generator @@ -23,6 +24,7 @@ def func(inserted): else: message = 'Field {} cannot be null'.format(field) raise SeederException(message) + return func def guess_field_formatters(self, faker): @@ -32,11 +34,10 @@ def guess_field_formatters(self, faker): :param faker: Faker factory object """ formatters = {} - model = self.model name_guesser = NameGuesser(faker) field_type_guesser = FieldTypeGuesser(faker) - for field in model._meta.fields: + for field in self.model._meta.fields: field_name = field.name if isinstance(field, (ForeignKey, ManyToManyField, OneToOneField)): formatters[field_name] = self.build_relation(field, field.rel.to) @@ -64,22 +65,31 @@ def execute(self, using, inserted_entities): :param using: :param inserted_entities: """ - obj = self.model() - - for field, format in self.field_formatters.items(): - if format: - if hasattr(format,'__call__'): - setattr(obj, field, format(inserted_entities)) - else: - setattr(obj, field, format) - obj.save(using=using) + def format_field(format, inserted_entities): + if callable(format): + return format(inserted_entities) + return format + + def turn_off_auto_add(model): + for field in model._meta.fields: + if getattr(field, 'auto_now', False): + field.auto_now = False + if getattr(field, 'auto_now_add', False): + field.auto_now_add = False + + manager = self.model.objects.db_manager(using=using) + turn_off_auto_add(manager.model) + + obj = manager.create(**{ + field: format_field(field_format, inserted_entities) + for field, field_format in self.field_formatters.items() + }) return obj.pk class Seeder(object): - def __init__(self, faker): """ :param faker: Generator @@ -89,7 +99,6 @@ def __init__(self, faker): self.quantities = {} self.orders = [] - def add_entity(self, model, number, customFieldFormatters=None): """ Add an order for the generation of $number records for $entity. @@ -130,7 +139,7 @@ def execute(self, using=None): number = self.quantities[klass] if klass not in inserted_entities: inserted_entities[klass] = [] - for i in range(0,number): + for i in range(0, number): entity = self.entities[klass].execute(using, inserted_entities) inserted_entities[klass].append(entity) diff --git a/django_seed/tests.py b/django_seed/tests.py index 64d264f..5f28f86 100644 --- a/django_seed/tests.py +++ b/django_seed/tests.py @@ -1,3 +1,4 @@ +from datetime import datetime from faker import Faker from django_seed.seeder import Seeder from django_seed.exceptions import SeederException, SeederCommandError @@ -21,7 +22,8 @@ class Game(models.Model): title = models.CharField(max_length=200) slug = models.SlugField(max_length=200) description = models.TextField() - created_at = models.DateTimeField() + created_at = models.DateTimeField(auto_now_add=True) + updated_at = models.DateTimeField(auto_now=True) updated_date = models.DateField() updated_time = models.TimeField() active = models.BooleanField() @@ -146,6 +148,30 @@ def test_no_entities_added(self): except Exception as e: self.assertTrue(isinstance(e, SeederException)) + def test_auto_now_add(self): + date = datetime(1957, 3, 6, 13, 13) + faker = fake + seeder = Seeder(faker) + seeder.add_entity(Game, 10, { + 'created_at': lambda x: date + }) + inserted_pks = seeder.execute()[Game] + + games = Game.objects.filter(pk__in=inserted_pks) + self.assertTrue(all(game.created_at == date for game in games)) + + def test_auto_now(self): + date = datetime(1957, 3, 6, 13, 13) + faker = fake + seeder = Seeder(faker) + seeder.add_entity(Game, 10, { + 'updated_at': lambda x: date + }) + inserted_pks = seeder.execute()[Game] + + games = Game.objects.filter(pk__in=inserted_pks) + self.assertTrue(all(game.updated_at == date for game in games)) + class APISeedTestCase(TestCase):