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

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

2 files changed, 12 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