From 385777c09a4f70402f10daaf79e0aa602e4e4071 Mon Sep 17 00:00:00 2001
From: Lars Wirzenius <liw@liw.fi>
Date: Sat, 16 May 2020 10:19:39 +0300
Subject: feat(template.py): call cleanup functions if defined

If a binding has a cleanup function defined, call it at the end of the
scenario, even if a step in the scenario failed.
---
 templates/python/template.py | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/templates/python/template.py b/templates/python/template.py
index 86396ff..76d8239 100644
--- a/templates/python/template.py
+++ b/templates/python/template.py
@@ -90,6 +90,7 @@ class Step:
         self._text = None
         self._args = {}
         self._function = None
+        self._cleanup = None
 
     def set_kind(self, kind):
         self._kind = kind
@@ -103,11 +104,22 @@ class Step:
     def set_function(self, function):
         self._function = function
 
+    def set_cleanup(self, cleanup):
+        self._cleanup = cleanup
+
     def do(self, ctx):
         print('  step: {} {}'.format(self._kind, self._text))
         logging.info('  step: {} {}'.format(self._kind, self._text))
         self._function(ctx, **self._args)
 
+    def cleanup(self, ctx):
+        if self._cleanup:
+            print('  cleanup: {} {}'.format(self._kind, self._text))
+            logging.info('  cleanup: {} {}'.format(self._kind, self._text))
+            self._cleanup(ctx)
+        else:
+            logging.info('  no cleanup defined: {} {}'.format(self._kind, self._text))
+
 
 class Scenario:
 
@@ -132,8 +144,18 @@ class Scenario:
         os.chdir(scendir)
 
         ctx = Context()
-        for step in self._steps:
-            step.do(ctx)
+        done = []
+        try:
+            for step in self._steps:
+                step.do(ctx)
+                done.append(step)
+        except Exception as e:
+            logging.error(str(e), exc_info=True)
+            for step in reversed(done):
+                step.cleanup(ctx)
+            raise
+        for step in reversed(done):
+            step.cleanup(ctx)
 
 
 {% for scenario in scenarios %}
@@ -149,6 +171,8 @@ class Scenario_{{ loop.index }}():
         step.set_kind('{{ step.kind | lower }}')
         step.set_text(decode_str('{{ step.text | base64 }}'))
         step.set_function({{ step.function }})
+        if '{{ step.cleanup }}':
+            step.set_cleanup({{ step.cleanup }})
         self._scenario.append_step(step)
         {% for part in step.parts %}{% if part.CapturedText is defined -%}
         name = decode_str('{{ part.CapturedText.name | base64 }}')
-- 
cgit v1.2.1