Merge pull request #256 from koterpillar/cyclic-definition-rescue

Don't leave AttributeBuilder in an inconsistent state on exceptions

2 files changed, 33 insertions, 9 deletions

diff --git a/factory/containers.py b/factory/containers.py
index d3f39c4..4961115 100644
--- a/factory/containers.py
+++ b/factory/containers.py
@@ -99,8 +99,10 @@ class LazyStub(object):
             val = self.__attrs[name]
             if isinstance(val, LazyValue):
                 self.__pending.append(name)
-                val = val.evaluate(self, self.__containers)
-                last = self.__pending.pop()
+                try:
+                    val = val.evaluate(self, self.__containers)
+                finally:
+                    last = self.__pending.pop()
                 assert name == last
             self.__values[name] = val
             return val
diff --git a/tests/test_containers.py b/tests/test_containers.py
index a308353..ea6ddff 100644
--- a/tests/test_containers.py
+++ b/tests/test_containers.py
@@ -79,18 +79,40 @@ class LazyStubTestCase(unittest.TestCase):
         self.assertEqual(2, stub.factory_parent.rank)
         self.assertEqual(1, stub.factory_parent.factory_parent.rank)
 
-    def test_cyclic_definition(self):
-        class LazyAttr(containers.LazyValue):
-            def __init__(self, attrname):
-                self.attrname = attrname
+    class LazyAttr(containers.LazyValue):
+        def __init__(self, attrname):
+            self.attrname = attrname
 
-            def evaluate(self, obj, container=None):
-                return 1 + getattr(obj, self.attrname)
+        def evaluate(self, obj, container=None):
+            return 1 + getattr(obj, self.attrname)
 
-        stub = containers.LazyStub({'one': LazyAttr('two'), 'two': LazyAttr('one')})
+    def test_cyclic_definition(self):
+        stub = containers.LazyStub({
+            'one': self.LazyAttr('two'),
+            'two': self.LazyAttr('one'),
+        })
 
         self.assertRaises(errors.CyclicDefinitionError, getattr, stub, 'one')
 
+    def test_cyclic_definition_rescue(self):
+        class LazyAttrDefault(self.LazyAttr):
+            def __init__(self, attname, defvalue):
+                super(LazyAttrDefault, self).__init__(attname)
+                self.defvalue = defvalue
+            def evaluate(self, obj, container=None):
+                try:
+                    return super(LazyAttrDefault, self).evaluate(obj, container)
+                except containers.CyclicDefinitionError:
+                    return self.defvalue
+
+        stub = containers.LazyStub({
+            'one': LazyAttrDefault('two', 10),
+            'two': self.LazyAttr('one'),
+        })
+
+        self.assertEqual(10, stub.one)
+        self.assertEqual(11, stub.two)
+
     def test_representation(self):
         class RandomObj(object):
             pass