1 files changed, 87 insertions, 32 deletions

diff --git a/templates/python/template.py b/templates/python/template.py
index 62b1fc9..86396ff 100644
--- a/templates/python/template.py
+++ b/templates/python/template.py
@@ -82,36 +82,91 @@ os.chdir(_datadir)
 #############################################################################
 # Code to implement the scenarios.
 
+
+class Step:
+
+    def __init__(self):
+        self._kind = None
+        self._text = None
+        self._args = {}
+        self._function = None
+
+    def set_kind(self, kind):
+        self._kind = kind
+
+    def set_text(self, text):
+        self._text = text
+
+    def set_arg(self, name, value):
+        self._args[name] = value
+
+    def set_function(self, function):
+        self._function = function
+
+    def do(self, ctx):
+        print('  step: {} {}'.format(self._kind, self._text))
+        logging.info('  step: {} {}'.format(self._kind, self._text))
+        self._function(ctx, **self._args)
+
+
+class Scenario:
+
+    def __init__(self):
+        self._title = None
+        self._steps = []
+
+    def get_title(self):
+        return self._title
+
+    def set_title(self, title):
+        self._title = title
+
+    def append_step(self, step):
+        self._steps.append(step)
+
+    def run(self):
+        print('scenario: {}'.format(self._title))
+        logging.info("Scenario: {}".format(self._title))
+        
+        scendir = tempfile.mkdtemp(dir=_datadir)
+        os.chdir(scendir)
+
+        ctx = Context()
+        for step in self._steps:
+            step.do(ctx)
+
+
 {% for scenario in scenarios %}
 ######################################
 # Scenario: {{ scenario.title }}
-def scenario_{{ loop.index }}():
-    title = decode_str('{{ scenario.title | base64 }}')
-    print('scenario: {}'.format(title))
-    logging.info("Scenario: {}".format(title))
-    _scendir = tempfile.mkdtemp(dir=_datadir)
-    os.chdir(_scendir)
-    ctx = Context()
-    {% for step in scenario.steps %}
-    # Step: {{ step.text }}
-    step = decode_str('{{ step.text | base64 }}')
-    print('  step: {{ step.kind | lower }} {}'.format(step))
-    logging.info('  step: {{ step.kind | lower }} {}'.format(step))
-    args = {}
-    {% for part in step.parts %}{% if part.CapturedText is defined -%}
-    name = decode_str('{{ part.CapturedText.name | base64 }}')
-    text = decode_str('{{ part.CapturedText.text | base64 }}')
-    args[name] = text
-    {% endif -%}
-    {% endfor -%}
-    {{ step.function }}(ctx, **args)
+class Scenario_{{ loop.index }}():
+    def __init__(self):
+        self._scenario = Scenario()
+        self._scenario.set_title(decode_str('{{ scenario.title | base64 }}'))
+        {% for step in scenario.steps %}
+        # Step: {{ step.text }}
+        step = Step()
+        step.set_kind('{{ step.kind | lower }}')
+        step.set_text(decode_str('{{ step.text | base64 }}'))
+        step.set_function({{ step.function }})
+        self._scenario.append_step(step)
+        {% for part in step.parts %}{% if part.CapturedText is defined -%}
+        name = decode_str('{{ part.CapturedText.name | base64 }}')
+        text = decode_str('{{ part.CapturedText.text | base64 }}')
+        step.set_arg(name, text)
+        {% endif -%}
+        {% endfor -%}
     {% endfor %}
-{% endfor %}
 
-_scenarios = {
-{% for scenario in scenarios %}
-    '{{ scenario.title }}': scenario_{{ loop.index }},
+    def get_title(self):
+        return self._scenario.get_title()
+
+    def run(self):
+        self._scenario.run()
 {% endfor %}
+
+_scenarios = { {% for scenario in scenarios %}
+    Scenario_{{ loop.index }}(),{% endfor %}
 }
 
 
@@ -156,20 +211,20 @@ def main():
     logging.info("patterns: {}".format(args.patterns))
     if len(args.patterns) == 0:
         logging.info("Executing all scenarios")
-        funcs = list(_scenarios.values())
-        random.shuffle(funcs)
+        todo = list(_scenarios)
+        random.shuffle(todo)
     else:
         logging.info("Executing requested scenarios only: {}".format(args.patterns))
         patterns = [arg.lower() for arg in args.patterns]
-        funcs = [
-            func
-            for title, func in _scenarios.items()
-            if any(pattern in title.lower() for pattern in patterns)
+        todo = [
+            scen
+            for scen in _scenarios
+            if any(pattern in scen.get_title().lower() for pattern in patterns)
         ]
 
     try:
-        for func in funcs:
-            func()
+        for scen in todo:
+            scen.run()
     except Exception as e:
         logging.error(str(e), exc_info=True)
         if args.save_on_failure: