Merge pull request #34 from jgadelange/master

Make compatible with auto_now(_add)

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)
 

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):