From a5184785b1229e02ec75506db675a2377b32c297 Mon Sep 17 00:00:00 2001 From: Raphaël Barrois Date: Wed, 14 Nov 2012 23:53:26 +0100 Subject: Keep FACTORY_FOR around. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit And add a test too. Signed-off-by: Raphaël Barrois --- factory/base.py | 3 --- tests/test_base.py | 8 ++++++++ 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 -- cgit v1.2.3