1 files changed, 32 insertions, 10 deletions
diff --git a/tests/util.py b/tests/util.py
index 2cd2c031..4bb6a653 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -19,7 +19,8 @@ except ImportError:
     # functools is new in 2.4
     wraps = lambda f: (lambda w: w)
 
-from sphinx import application, builder
+from sphinx import application
+from sphinx.ext.autodoc import AutoDirective
 
 from path import path
 
@@ -29,7 +30,7 @@ from nose import tools
 __all__ = [
     'test_root',
     'raises', 'raises_msg', 'Struct',
-    'ListOutput', 'TestApp', 'with_app',
+    'ListOutput', 'TestApp', 'with_app', 'gen_with_app',
     'path', 'with_tempdir', 'write_file',
     'sprint',
 ]
@@ -97,12 +98,14 @@ class TestApp(application.Sphinx):
     """
 
     def __init__(self, srcdir=None, confdir=None, outdir=None, doctreedir=None,
-                 buildername='html', confoverrides=None, status=None, warning=None,
-                 freshenv=None, confname='conf.py', cleanenv=False):
+                 buildername='html', confoverrides=None,
+                 status=None, warning=None, freshenv=None,
+                 warningiserror=None, tags=None,
+                 confname='conf.py', cleanenv=False):
 
         application.CONFIG_FILENAME = confname
 
-        self.cleanup_trees = []
+        self.cleanup_trees = [test_root / 'generated']
 
         if srcdir is None:
             srcdir = test_root
@@ -134,12 +137,15 @@ class TestApp(application.Sphinx):
             warning = ListOutput('stderr')
         if freshenv is None:
             freshenv = False
+        if warningiserror is None:
+            warningiserror = False
 
         application.Sphinx.__init__(self, srcdir, confdir, outdir, doctreedir,
                                     buildername, confoverrides, status, warning,
-                                    freshenv)
+                                    freshenv, warningiserror, tags)
 
     def cleanup(self, doctrees=False):
+        AutoDirective._registry.clear()
         for tree in self.cleanup_trees:
             shutil.rmtree(tree, True)
 
@@ -153,10 +159,26 @@ def with_app(*args, **kwargs):
         @wraps(func)
         def deco(*args2, **kwargs2):
             app = TestApp(*args, **kwargs)
-            try:
-                func(app, *args2, **kwargs2)
-            finally:
-                app.cleanup()
+            func(app, *args2, **kwargs2)
+            # don't execute cleanup if test failed
+            app.cleanup()
+        return deco
+    return generator
+
+
+def gen_with_app(*args, **kwargs):
+    """
+    Make a TestApp with args and kwargs, pass it to the test and clean up
+    properly.
+    """
+    def generator(func):
+        @wraps(func)
+        def deco(*args2, **kwargs2):
+            app = TestApp(*args, **kwargs)
+            for item in func(app, *args2, **kwargs2):
+                yield item
+            # don't execute cleanup if test failed
+            app.cleanup()
         return deco
     return generator