Keep FACTORY_FOR around.

And add a test too.

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

2 files changed, 8 insertions, 3 deletions

diff --git a/factory/base.py b/factory/base.py
index 3d4e8c1..2b78e2b 100644
--- a/factory/base.py
+++ b/factory/base.py
@@ -236,9 +236,6 @@ class FactoryMetaClass(BaseFactoryMetaClass):
         associated_class = cls._discover_associated_class(class_name, attrs,
                 inherited_associated_class)
 
-        # Remove the FACTORY_CLASS_DECLARATION attribute from attrs, if present.
-        attrs.pop(FACTORY_CLASS_DECLARATION, None)
-
         # If inheriting the factory from a parent, keep a link to it.
         # This allows to use the sequence counters from the parents.
         if associated_class == inherited_associated_class:
diff --git a/tests/test_base.py b/tests/test_base.py
index 63fc62b..64570cd 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -64,6 +64,14 @@ class SafetyTestCase(unittest.TestCase):
 
 
 class FactoryTestCase(unittest.TestCase):
+    def test_factory_for(self):
+        class TestObjectFactory(base.Factory):
+            FACTORY_FOR = TestObject
+
+        self.assertEqual(TestObject, TestObjectFactory.FACTORY_FOR)
+        obj = TestObjectFactory.build()
+        self.assertFalse(hasattr(obj, 'FACTORY_FOR'))
+
     def testDisplay(self):
         class TestObjectFactory(base.Factory):
             FACTORY_FOR = FakeDjangoModel