yarnlib: Split out scenario and step setup

This is so the future scenario multi-runner can re-use the common setup
code.

1 files changed, 30 insertions, 17 deletions

diff --git a/yarnlib/scenario_runner.py b/yarnlib/scenario_runner.py
index cb74859..91e7f0f 100644
--- a/yarnlib/scenario_runner.py
+++ b/yarnlib/scenario_runner.py
@@ -114,34 +114,41 @@ class ScenarioRunner(object):
         self.cmdrunner = cmdrunner
         self.testdir = testdir
 
+    def setup_scenario(self, scenario):
+        scenario_dir = self.scenario_dir(self.testdir, scenario)
+        os.mkdir(scenario_dir)
+        datadir = self.datadir(scenario_dir)
+        os.mkdir(datadir)
+        homedir = self.homedir(datadir)
+        os.mkdir(homedir)
+
+        ud = self.pre_scenario_cb(scenario, datadir, homedir)
+        return scenario_dir, datadir, homedir, ud
+
     def run_scenarios(self, scenarios):
         failed = []
         for scenario in scenarios:
-            scenario_dir = self.scenario_dir(self.testdir, scenario)
-            os.mkdir(scenario_dir)
-            datadir = self.datadir(scenario_dir)
-            os.mkdir(datadir)
-            homedir = self.homedir(datadir)
-            os.mkdir(homedir)
-
-            ud = self.pre_scenario_cb(scenario, datadir, homedir)
+            scenario_dir, datadir, homedir, ud = self.setup_scenario(scenario)
             if not self.run_scenario(scenario, datadir, homedir):
                 failed.append(scenario)
             self.post_scenario_cb(scenario, datadir, homedir, ud)
         return failed
 
-    def run_scenario(self, scenario, datadir, homedir):
+    @staticmethod
+    def partition_steps(scenario):
         assuming = [s for s in scenario.steps if s.what == 'ASSUMING']
         cleanup = [s for s in scenario.steps if s.what == 'FINALLY']
-        normal = [s for s in scenario.steps if s not in assuming + cleanup]
+        normal = [s for s in scenario.steps
+                  if s.what not in ('ASSUMING', 'FINALLY')]
+        return assuming, cleanup, normal
+
+    def run_scenario(self, scenario, datadir, homedir):
+        assuming, cleanup, normal = self.partition_steps(scenario)
+        scenario_env = dict(self.env, HOME=homedir, DATADIR=datadir)
 
         ok = True
         step_number = 1
 
-        scenario_env = dict(self.env)
-        scenario_env['HOME'] = homedir
-        scenario_env['DATADIR'] = datadir
-
         for step in assuming:
             exit = self.run_step(scenario, step, scenario_env, step_number)
             step_number += 1
@@ -164,12 +171,12 @@ class ScenarioRunner(object):
 
         return ok
 
-    def run_step(self, scenario, step, scenario_env, step_number):
+    def setup_step(self, step, scenario_env, scenario, step_number):
         m = yarnlib.implements_matches_step(step.implementation, step)
         assert m is not None
         step_env = dict(scenario_env)
-        for i, match in enumerate(m.groups('')):
-            step_env['MATCH_%d' % (i+1)] = match
+        step_env.update(('MATCH_%d' % i, match) for (i, match)
+                        in enumerate(m.groups(''), 1))
 
         # All parameters passed as keyword-arguments, so that the callback
         # may declare parameters in any order, and ignore any parameters
@@ -180,6 +187,12 @@ class ScenarioRunner(object):
 
         shell_script = '%s\n\n%s\n' % (
             self.shell_prelude, step.implementation.shell)
+        return step_env, pre_step_userdata, shell_script
+
+    def run_step(self, scenario, step, scenario_env, step_number):
+        step_env, pre_step_userdata, shell_script = self.setup_step(
+                        step=step, scenario_env=scenario_env,
+                        scenario=scenario, step_number=step_number)
         exit, stdout, stderr = self.cmdrunner(
             ['sh', '-xeuc', shell_script], env=step_env, cwd=self.srcdir)