Add batches.

Signed-off-by: Raphaël Barrois <raphael.barrois@polytechnique.org>

3 files changed, 125 insertions, 1 deletions

diff --git a/factory/__init__.py b/factory/__init__.py
index e58dc51..4acd380 100644
--- a/factory/__init__.py
+++ b/factory/__init__.py
@@ -33,6 +33,10 @@ from base import (
     stub,
     make_factory,
 
+    build_batch,
+    create_batch,
+    stub_batch,
+
     BUILD_STRATEGY,
     CREATE_STRATEGY,
     STUB_STRATEGY,
diff --git a/factory/base.py b/factory/base.py
index c17b7ce..99cf49d 100644
--- a/factory/base.py
+++ b/factory/base.py
@@ -309,11 +309,35 @@ class BaseFactory(object):
         raise cls.UnsupportedStrategy()
 
     @classmethod
+    def build_batch(cls, size, **kwargs):
+        """Build a batch of instances of the given class, with overriden attrs.
+
+        Args:
+            size (int): the number of instances to build
+        
+        Returns:
+            object list: the built instances
+        """
+        return [cls.build(**kwargs) for _ in xrange(size)]
+
+    @classmethod
     def create(cls, **kwargs):
         """Create an instance of the associated class, with overriden attrs."""
         raise cls.UnsupportedStrategy()
 
     @classmethod
+    def create_batch(cls, size, **kwargs):
+        """Create a batch of instances of the given class, with overriden attrs.
+
+        Args:
+            size (int): the number of instances to create
+        
+        Returns:
+            object list: the created instances
+        """
+        return [cls.create(**kwargs) for _ in xrange(size)]
+
+    @classmethod
     def stub(cls, **kwargs):
         """Retrieve a stub of the associated class, with overriden attrs.
 
@@ -325,6 +349,18 @@ class BaseFactory(object):
             setattr(stub_object, name, value)
         return stub_object
 
+    @classmethod
+    def stub_batch(cls, size, **kwargs):
+        """Stub a batch of instances of the given class, with overriden attrs.
+
+        Args:
+            size (int): the number of instances to stub
+        
+        Returns:
+            object list: the stubbed instances
+        """
+        return [cls.stub(**kwargs) for _ in xrange(size)]
+
 
 class StubFactory(BaseFactory):
     __metaclass__ = BaseFactoryMetaClass
@@ -468,10 +504,28 @@ def build(klass, **kwargs):
     """Create a factory for the given class, and build an instance."""
     return make_factory(klass, **kwargs).build()
 
+
+def build_batch(klass, size, **kwargs):
+    """Create a factory for the given class, and build a batch of instances."""
+    return make_factory(klass, **kwargs).build_batch(size)
+
+
 def create(klass, **kwargs):
     """Create a factory for the given class, and create an instance."""
     return make_factory(klass, **kwargs).create()
 
+
+def create_batch(klass, size, **kwargs):
+    """Create a factory for the given class, and create a batch of instances."""
+    return make_factory(klass, **kwargs).create_batch(size)
+
+
 def stub(klass, **kwargs):
     """Create a factory for the given class, and stub an instance."""
     return make_factory(klass, **kwargs).stub()
+
+
+def stub_batch(klass, size, **kwargs):
+    """Create a factory for the given class, and stub a batch of instances."""
+    return make_factory(klass, **kwargs).stub_batch(size)
+
diff --git a/tests/test_using.py b/tests/test_using.py
index 85a12ca..7ed94e2 100644
--- a/tests/test_using.py
+++ b/tests/test_using.py
@@ -69,11 +69,34 @@ class SimpleBuildTestCase(unittest.TestCase):
         self.assertEqual(obj.three, 3)
         self.assertEqual(obj.four, None)
 
+    def test_build_batch(self):
+        objs = factory.build_batch(TestObject, 4, two=2,
+            three=factory.LazyAttribute(lambda o: o.two + 1))
+
+        self.assertEqual(4, len(objs))
+        self.assertEqual(4, len(set(objs)))
+
+        for obj in objs:
+            self.assertEqual(obj.one, None)
+            self.assertEqual(obj.two, 2)
+            self.assertEqual(obj.three, 3)
+            self.assertEqual(obj.four, None)
+
     def test_create(self):
         obj = factory.create(FakeDjangoModel, foo='bar')
         self.assertEqual(obj.id, 1)
         self.assertEqual(obj.foo, 'bar')
 
+    def test_create_batch(self):
+        objs = factory.create_batch(FakeDjangoModel, 4, foo='bar')
+
+        self.assertEqual(4, len(objs))
+        self.assertEqual(4, len(set(objs)))
+
+        for obj in objs:
+            self.assertEqual(obj.id, 1)
+            self.assertEqual(obj.foo, 'bar')
+
     def test_stub(self):
         obj = factory.stub(TestObject, three=3)
         self.assertEqual(obj.three, 3)
@@ -133,6 +156,19 @@ class FactoryTestCase(unittest.TestCase):
         self.assertEqual('one43', test_object1.one)
         self.assertEqual('two43', test_object1.two)
 
+    def test_sequence_batch(self):
+        class TestObjectFactory(factory.Factory):
+            one = factory.Sequence(lambda n: 'one' + n)
+            two = factory.Sequence(lambda n: 'two' + n)
+
+        objs = TestObjectFactory.build_batch(20)
+
+        self.assertEqual(20, len(objs))
+        self.assertEqual(20, len(set(objs)))
+        for i, obj in enumerate(objs):
+            self.assertEqual('one%d' % i, obj.one)
+            self.assertEqual('two%d' % i, obj.two)
+
     def testLazyAttribute(self):
         class TestObjectFactory(factory.Factory):
             one = factory.LazyAttribute(lambda a: 'abc' )
@@ -225,6 +261,37 @@ class FactoryTestCase(unittest.TestCase):
         self.assertEqual(test_model.one, 'one')
         self.assertTrue(test_model.id)
 
+    def test_create_batch(self):
+        class TestModelFactory(factory.Factory):
+            one = 'one'
+
+        objs = TestModelFactory.create_batch(20, two=factory.Sequence(int))
+
+        self.assertEqual(20, len(objs))
+        self.assertEqual(20, len(set(objs)))
+
+        for i, obj in enumerate(objs):
+            self.assertEqual('one', obj.one)
+            self.assertEqual(i, obj.two)
+            self.assertTrue(obj.id)
+
+    def test_stub_batch(self):
+        class TestObjectFactory(factory.Factory):
+            one = 'one'
+            two = factory.LazyAttribute(lambda a: a.one + ' two')
+            three = factory.Sequence(lambda n: int(n))
+
+        objs = TestObjectFactory.stub_batch(20,
+            one=factory.Sequence(lambda n: n))
+
+        self.assertEqual(20, len(objs))
+        self.assertEqual(20, len(set(objs)))
+
+        for i, obj in enumerate(objs):
+            self.assertEqual(str(i), obj.one)
+            self.assertEqual('%d two' % i, obj.two)
+            self.assertEqual(i, obj.three)
+
     def testInheritance(self):
         class TestObjectFactory(factory.Factory):
             one = 'one'
@@ -364,7 +431,6 @@ class SubFactoryTestCase(unittest.TestCase):
         self.assertEqual(wrapping.wrapped.three, 3)
         self.assertEqual(wrapping.wrapped.four, 4)
 
-
     def testNestedSubFactory(self):
         """Test nested sub-factories."""