Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make compatible with auto_now(_add) #34

Merged
merged 2 commits into from
Mar 31, 2016
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
45 changes: 27 additions & 18 deletions django_seed/seeder.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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):
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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)

Expand Down
28 changes: 27 additions & 1 deletion django_seed/tests.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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):

Expand Down