Add a '@use_strategy' decorator for forcing alternate strategies.

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

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):