Add support for 'LazyContainerAttribute', when a SubFactory field needs access to attributes from its parent.

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

2 files changed, 80 insertions, 1 deletions

diff --git a/tests/test_base.py b/tests/test_base.py
index aaaa65f..9e61378 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -467,6 +467,60 @@ class FactoryDefaultStrategyTestCase(unittest.TestCase):
         self.assertEqual(wrapping.wrapped.two, 4)
         self.assertEqual(wrapping.friend, 5)
 
+    def testDiamondSubFactory(self):
+        """Tests the case where an object has two fields with a common field."""
+        class InnerMost(object):
+            def __init__(self, a, b):
+                self.a = a
+                self.b = b
+
+        class SideA(object):
+            def __init__(self, inner_from_a):
+                self.inner_from_a = inner_from_a
+
+        class SideB(object):
+            def __init__(self, inner_from_b):
+                self.inner_from_b = inner_from_b
+
+        class OuterMost(object):
+            def __init__(self, foo, side_a, side_b):
+                self.foo = foo
+                self.side_a = side_a
+                self.side_b = side_b
+
+        class InnerMostFactory(base.Factory):
+            FACTORY_FOR = InnerMost
+            a = 15
+            b = 20
+
+        class SideAFactory(base.Factory):
+            FACTORY_FOR = SideA
+            inner_from_a = declarations.SubFactory(InnerMostFactory, a=20)
+
+        class SideBFactory(base.Factory):
+            FACTORY_FOR = SideB
+            inner_from_b = declarations.SubFactory(InnerMostFactory, b=15)
+
+        class OuterMostFactory(base.Factory):
+            FACTORY_FOR = OuterMost
+
+            foo = 30
+            side_a = declarations.SubFactory(SideAFactory,
+                inner_from_a__a=declarations.LazyContainerAttribute(lambda obj, containers: containers[1].foo * 2))
+            side_b = declarations.SubFactory(SideBFactory,
+                inner_from_b=declarations.LazyContainerAttribute(lambda obj, containers: containers[0].side_a.inner_from_a))
+
+        outer = OuterMostFactory.build()
+        self.assertEqual(outer.foo, 30)
+        self.assertEqual(outer.side_a.inner_from_a, outer.side_b.inner_from_b)
+        self.assertEqual(outer.side_a.inner_from_a.a, outer.foo * 2)
+        self.assertEqual(outer.side_a.inner_from_a.b, 20)
+
+        outer = OuterMostFactory.build(side_a__inner_from_a__b = 4)
+        self.assertEqual(outer.foo, 30)
+        self.assertEqual(outer.side_a.inner_from_a, outer.side_b.inner_from_b)
+        self.assertEqual(outer.side_a.inner_from_a.a, outer.foo * 2)
+        self.assertEqual(outer.side_a.inner_from_a.b, 4)
 
     def testStubStrategy(self):
         base.Factory.default_strategy = base.STUB_STRATEGY
diff --git a/tests/test_containers.py b/tests/test_containers.py
index 34b61d4..57c06cf 100644
--- a/tests/test_containers.py
+++ b/tests/test_containers.py
@@ -43,12 +43,37 @@ class LazyStubTestCase(unittest.TestCase):
 
         self.assertRaises(AttributeError, getattr, stub, 'three')
 
+    def test_accessing_container(self):
+        class LazyAttr(containers.LazyValue):
+            def __init__(self, obj_attr, container_attr):
+                self.obj_attr = obj_attr
+                self.container_attr = container_attr
+
+            def evaluate(self, obj, containers=()):
+                if containers:
+                    add = getattr(containers[0], self.container_attr)
+                else:
+                    add = 0
+                return getattr(obj, self.obj_attr) + add
+
+        class DummyContainer(object):
+            three = 3
+
+        stub = containers.LazyStub({'one': LazyAttr('two', 'three'), 'two': 2, 'three': 42},
+            containers=(DummyContainer(),))
+
+        self.assertEqual(5, stub.one)
+
+        stub = containers.LazyStub({'one': LazyAttr('two', 'three'), 'two': 2, 'three': 42},
+            containers=())
+        self.assertEqual(2, stub.one)
+
     def test_cyclic_definition(self):
         class LazyAttr(containers.LazyValue):
             def __init__(self, attrname):
                 self.attrname = attrname
 
-            def evaluate(self, obj):
+            def evaluate(self, obj, container=None):
                 return 1 + getattr(obj, self.attrname)
 
         stub = containers.LazyStub({'one': LazyAttr('two'), 'two': LazyAttr('one')})