diff options
| -rw-r--r-- | factory/__init__.py | 1 | ||||
| -rw-r--r-- | factory/base.py | 11 | ||||
| -rw-r--r-- | tests/test_base.py | 8 |
3 files changed, 20 insertions, 0 deletions
diff --git a/factory/__init__.py b/factory/__init__.py index 1d4408f..dd91343 100644 --- a/factory/__init__.py +++ b/factory/__init__.py @@ -44,6 +44,7 @@ from base import ( BUILD_STRATEGY, CREATE_STRATEGY, STUB_STRATEGY, + use_strategy, DJANGO_CREATION, NAIVE_BUILD, diff --git a/factory/base.py b/factory/base.py index 62131fb..ef782c7 100644 --- a/factory/base.py +++ b/factory/base.py @@ -614,3 +614,14 @@ def simple_generate(klass, create, **kwargs): def simple_generate_batch(klass, create, size, **kwargs): """Create a factory for the given class, and simple_generate instances.""" return make_factory(klass, **kwargs).simple_generate_batch(create, size) + + +def use_strategy(new_strategy): + """Force the use of a different strategy. + + This is an alternative to setting default_strategy in the class definition. + """ + def wrapped_class(klass): + klass.default_strategy = new_strategy + return klass + return wrapped_class diff --git a/tests/test_base.py b/tests/test_base.py index 8da655e..a548b56 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -141,6 +141,14 @@ class FactoryDefaultStrategyTestCase(unittest.TestCase): TestModelFactory.default_strategy = base.BUILD_STRATEGY self.assertRaises(base.StubFactory.UnsupportedStrategy, TestModelFactory) + def test_change_strategy(self): + @base.use_strategy(base.CREATE_STRATEGY) + class TestModelFactory(base.StubFactory): + one = 'one' + + self.assertEqual(base.CREATE_STRATEGY, TestModelFactory.default_strategy) + + class FactoryCreationTestCase(unittest.TestCase): def testFactoryFor(self): class TestFactory(base.Factory): |