[Telemetry] Add Benchmark.ShouldTearDownStateAfterEachStorySetRun

telemetry/telemetry/internal/story_runner.py

Index: telemetry/telemetry/internal/story_runner.py
diff --git a/telemetry/telemetry/internal/story_runner.py b/telemetry/telemetry/internal/story_runner.py
index 5e6a32fb711ebb271f53c1f8b1410682325fb8d5..cd0c9865b243e90a8b3916248d6c487ce7e780d7 100644
--- a/telemetry/telemetry/internal/story_runner.py
+++ b/telemetry/telemetry/internal/story_runner.py
@@ -172,7 +172,7 @@ def StoriesGroupedByStateClass(story_set, allow_multiple_groups):
 
 
 def Run(test, story_set, finder_options, results, max_failures=None,
-        should_tear_down_state_after_each_story_run=False):
+        tear_down_after_story=False, tear_down_after_story_set=False):
   """Runs a given test against a given page_set with the given options.
 
   Stop execution for unexpected exceptions such as KeyboardInterrupt.
@@ -247,13 +247,16 @@ def Run(test, story_set, finder_options, results, max_failures=None,
                 # Print current exception and propagate existing exception.
                 exception_formatter.PrintFormattedException(
                     msg='Exception from result processing:')
-              if state and should_tear_down_state_after_each_story_run:
+              if state and tear_down_after_story:
                 state.TearDownState()
                 state = None
           if (effective_max_failures is not None and
               len(results.failures) > effective_max_failures):
             logging.error('Too many failures. Aborting.')
             return
+        if state and tear_down_after_story_set:
+          state.TearDownState()
+          state = None
     finally:
       if state:
         has_existing_exception = sys.exc_info() != (None, None, None)
@@ -310,7 +313,8 @@ def RunBenchmark(benchmark, finder_options):
       benchmark.ValueCanBeAddedPredicate) as results:
     try:
       Run(pt, stories, finder_options, results, benchmark.max_failures,
-          benchmark.ShouldTearDownStateAfterEachStoryRun())
+          benchmark.ShouldTearDownStateAfterEachStoryRun(),
+          benchmark.ShouldTearDownStateAfterEachStorySetRun())
       return_code = min(254, len(results.failures))
     except Exception:
       exception_formatter.PrintFormattedException()