Fixed multiple add entity bug and added tests

django_seed/seeder.py

@@ -109,8 +109,6 @@ def __init__(self, faker):
         :param faker: Generator
         """
         self.faker = faker
-        self.entities = {}
-        self.quantities = {}
         self.orders = []
 
     def add_entity(self, model, number, customFieldFormatters=None):
@@ -126,17 +124,21 @@ def add_entity(self, model, number, customFieldFormatters=None):
         callable as value
         :type customFieldFormatters: dict or None
         """
-        if not isinstance(model, ModelSeeder):
-            model = ModelSeeder(model)
+
+        # We always want to make a new ModelSeeder in case multiple unique
+        # orders for a specific model are created before a single execute
+        model = ModelSeeder(model)
 
         model.field_formatters = model.guess_field_formatters(self.faker)
         if customFieldFormatters:
             model.field_formatters.update(customFieldFormatters)
 
-        klass = model.model
-        self.entities[klass] = model
-        self.quantities[klass] = number
-        self.orders.append(klass)
+        order = {
+            "klass": model.model,
+            "quantity": number,
+            "entity": model,
+        }
+        self.orders.append(order)
 
     def execute(self, using=None):
         """
@@ -149,13 +151,17 @@ def execute(self, using=None):
             using = self.get_connection()
 
         inserted_entities = {}
-        for klass in self.orders:
-            number = self.quantities[klass]
+        while len(self.orders):
+            order = self.orders.pop(0)
+            number = order["quantity"]
+            klass = order["klass"]
+            entity = order["entity"]
+
             if klass not in inserted_entities:
                 inserted_entities[klass] = []
             for i in range(0, number):
-                entity = self.entities[klass].execute(using, inserted_entities)
-                inserted_entities[klass].append(entity)
+                executed_entity = entity.execute(using, inserted_entities)
+                inserted_entities[klass].append(executed_entity)
 
         return inserted_entities
 
@@ -165,10 +171,10 @@ def get_connection(self):
         :rtype: Connection
         """
 
-        klass = self.entities.keys()
-        if not klass:
+        klasses = [order["klass"] for order in self.orders]
+        if not klasses:
             message = 'No classed found. Did you add entities to the Seeder?'
             raise SeederException(message)
-        klass = list(klass)[0]
+        klass = list(klasses)[0]
 
         return klass.objects._db

django_seed/tests.py

@@ -147,6 +147,29 @@ def test_population(self):
         seeder = Seeder(faker)
         seeder.add_entity(Game, 10)
         self.assertEqual(len(seeder.execute()[Game]), 10)
+        self.assertEqual(len(Game.objects.all()), 10)
+
+        seeder.add_entity(Game, 40)
+        self.assertEqual(len(seeder.execute()[Game]), 40)
+        self.assertEqual(len(Game.objects.all()), 50)
+
+    def test_same_model_unique_fields(self):
+        faker = fake
+        seeder = Seeder(faker)
+        seeder.add_entity(Game, 10, {
+            "title": "First Game"
+        })
+
+        seeder.add_entity(Game, 20, {
+            "title": "Second Game"
+        })
+
+        inserted_pks = seeder.execute()
+
+        self.assertEqual(len(inserted_pks[Game]), 30)
+        self.assertEqual(len(Game.objects.all()), 30)
+        self.assertEqual(Game.objects.get(id=inserted_pks[Game][0]).title, "First Game")
+        self.assertEqual(Game.objects.get(id=inserted_pks[Game][-1]).title, "Second Game")
 
     def test_guesser(self):
         faker = fake