Add finalize_teardown method to clean up if setup or teardown fail

asv/benchmark.py

@@ -442,6 +442,7 @@ def __init__(self, name, func, attr_sources):
         self.pretty_name = getattr(func, "pretty_name", None)
         self._attr_sources = attr_sources
         self._setups = list(_get_all_attrs(attr_sources, 'setup', True))[::-1]
+        self._finalize_teardowns = list(_get_all_attrs(attr_sources, 'finalize_teardown', True))
         self._teardowns = list(_get_all_attrs(attr_sources, 'teardown', True))
         self._setup_cache = _get_first_attr(attr_sources, 'setup_cache', None)
         self.setup_cache_key = get_setup_cache_key(self._setup_cache)
@@ -590,6 +591,11 @@ def method_caller():
                 code, {'run': self.func, 'params': self._current_params},
                 {}, filename)
 
+    def do_finalize_teardown(self, exc_type, exception, trace):
+        """Called if the setup or the teardown of any test raise an exception.
+        """
+        for finalize_teardown in self._finalize_teardowns:
+            finalize_teardown(exc_type, exception, trace, self._current_params)
 
 class TimeBenchmark(Benchmark):
     """
@@ -1151,17 +1157,22 @@ def main_run(args):
         if cache is not None:
             benchmark.insert_param(cache)
 
-    skip = benchmark.do_setup()
-
     try:
-        if skip:
-            result = float('nan')
-        else:
-            result = benchmark.do_run()
-            if profile_path is not None:
-                benchmark.do_profile(profile_path)
+        skip = benchmark.do_setup()
+
+        try:
+            if skip:
+                result = float('nan')
+            else:
+                result = benchmark.do_run()
+                if profile_path is not None:
+                    benchmark.do_profile(profile_path)
+        finally:
+            benchmark.do_teardown()
     finally:
-        benchmark.do_teardown()
+        exc_type, exc, trace = sys.exc_info()
+        benchmark.do_finalize_teardown(exc_type, exc, trace)
+        raise exc
 
     # Write the output value
     with open(result_file, 'w') as fp:

docs/source/writing_benchmarks.rst

@@ -170,6 +170,29 @@ The ``setup_cache`` timeout can be specified by setting the
 ``.timeout`` attribute of the ``setup_cache`` function. The default
 value is the maximum of the timeouts of the benchmarks using it.
 
+If the ``setup`` or ``teardown`` of the test function raise an
+exception, the ``finalize_teardown`` method (if it exists) will be
+called with the exception as first argument, and the current
+benchmark parameters as second. The exception is then re-raised.
+
+This method gives you a chance to clean up any side-effects (or
+shared state) of the suite::
+
+  class Suite:
+      def setup(self, *args, **kwargs):
+          if not hasattr(self, '_server'):
+              self._server = MySuperServer()
+              self._server.launch()
+          # other code
+
+      def finalize_teardown(self, exc_info, current_params):
+          if hasattr(self, '_server'):
+              self._server.stop()
+
+Had the ``finalize_teardown`` not stopped the server, ASV would have
+waited until the ``setup`` timeout to go to the next benchmark, since
+a running process was still attached.
+
 .. _benchmark-attributes:
 
 Benchmark attributes

test/benchmark/finalize_teardown.py

@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+
+
+class FinalizeTeardownSetupTest:
+    def setup(self, n):
+        raise ValueError("Setup fails")
+
+    def time_it(self, n):
+        pass
+
+    def finalize_teardown(self, exc_type, exc, trace):
+        assert exc_type == ValueError
+        assert str(exc) == "Setup fails"
+        assert trace is not None
+
+
+class FinalizeTeardownTeardownTest:
+    def teardown(self, n):
+        raise NotImplementedError("Teardown fails")
+
+    def time_it(self, n):
+        pass
+
+    def finalize_teardown(self, exc_type, exc, trace):
+        assert exc_type == NotImplementedError
+        assert str(exc) == "Teardown fails"
+        assert trace is not None