diff options
-rw-r--r-- | factory/base.py | 3 | ||||
-rw-r--r-- | 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 |