Merge pull request #267 from bors-ltd/master

factory: LazyFunction to just call a function in the simplest case

2 files changed, 30 insertions, 4 deletions

diff --git a/tests/test_base.py b/tests/test_base.py
index 24f64e5..dd74e35 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -134,25 +134,28 @@ class OptionsTests(unittest.TestCase):
         self.assertEqual(AbstractFactory, AbstractFactory._meta.counter_reference)
 
     def test_declaration_collecting(self):
-        lazy = declarations.LazyAttribute(lambda _o: 1)
+        lazy = declarations.LazyFunction(int)
+        lazy2 = declarations.LazyAttribute(lambda _o: 1)
         postgen = declarations.PostGenerationDeclaration()
 
         class AbstractFactory(base.Factory):
             x = 1
             y = lazy
+            y2 = lazy2
             z = postgen
 
         # Declarations aren't removed
         self.assertEqual(1, AbstractFactory.x)
         self.assertEqual(lazy, AbstractFactory.y)
+        self.assertEqual(lazy2, AbstractFactory.y2)
         self.assertEqual(postgen, AbstractFactory.z)
 
         # And are available in class Meta
-        self.assertEqual({'x': 1, 'y': lazy}, AbstractFactory._meta.declarations)
+        self.assertEqual({'x': 1, 'y': lazy, 'y2': lazy2}, AbstractFactory._meta.declarations)
         self.assertEqual({'z': postgen}, AbstractFactory._meta.postgen_declarations)
 
     def test_inherited_declaration_collecting(self):
-        lazy = declarations.LazyAttribute(lambda _o: 1)
+        lazy = declarations.LazyFunction(int)
         lazy2 = declarations.LazyAttribute(lambda _o: 2)
         postgen = declarations.PostGenerationDeclaration()
         postgen2 = declarations.PostGenerationDeclaration()
@@ -178,7 +181,7 @@ class OptionsTests(unittest.TestCase):
         self.assertEqual({'z': postgen, 'b': postgen2}, OtherFactory._meta.postgen_declarations)
 
     def test_inherited_declaration_shadowing(self):
-        lazy = declarations.LazyAttribute(lambda _o: 1)
+        lazy = declarations.LazyFunction(int)
         lazy2 = declarations.LazyAttribute(lambda _o: 2)
         postgen = declarations.PostGenerationDeclaration()
         postgen2 = declarations.PostGenerationDeclaration()
diff --git a/tests/test_containers.py b/tests/test_containers.py
index 083b306..825e897 100644
--- a/tests/test_containers.py
+++ b/tests/test_containers.py
@@ -215,6 +215,29 @@ class AttributeBuilderTestCase(unittest.TestCase):
         ab = containers.AttributeBuilder(FakeFactory, extra={'one': seq2})
         self.assertEqual({'one': 'yy1'}, ab.build(create=False))
 
+    def test_lazy_function(self):
+        lf = declarations.LazyFunction(int)
+
+        class FakeFactory(object):
+            @classmethod
+            def declarations(cls, extra):
+                d = {'one': 1, 'two': lf}
+                d.update(extra)
+                return d
+
+            @classmethod
+            def _generate_next_sequence(cls):
+                return 1
+
+        ab = containers.AttributeBuilder(FakeFactory)
+        self.assertEqual({'one': 1, 'two': 0}, ab.build(create=False))
+
+        ab = containers.AttributeBuilder(FakeFactory, {'one': 4})
+        self.assertEqual({'one': 4, 'two': 0}, ab.build(create=False))
+
+        ab = containers.AttributeBuilder(FakeFactory, {'one': 4, 'three': lf})
+        self.assertEqual({'one': 4, 'two': 0, 'three': 0}, ab.build(create=False))
+
     def test_lazy_attribute(self):
         la = declarations.LazyAttribute(lambda a: a.one * 2)