refactor(templates/python): make helpers self-contained

To make every helper module be self-contained, move things around so
that a helper doesn't refer to global symbols from another helper.
This will allow us to use things like flake8 to statically check the
helper modules for common errors.

4 files changed, 37 insertions, 18 deletions

diff --git a/templates/python/files.py b/templates/python/files.py
index 9adbd30..6346172 100644
--- a/templates/python/files.py
+++ b/templates/python/files.py
@@ -1,3 +1,23 @@
 # Retrieve an embedded test data file using filename.
+
+
+class Files:
+    def __init__(self):
+        self._files = {}
+
+    def set(self, filename, content):
+        self._files[filename] = content
+
+    def get(self, filename):
+        return self._files[filename]
+
+
+_files = Files()
+
+
+def store_file(filename, content):
+    _files.set(filename, content)
+
+
 def get_file(filename):
-    return _files[filename]
+    return _files.get(filename)
diff --git a/templates/python/main.py b/templates/python/main.py
index 3a39164..fa78a5f 100644
--- a/templates/python/main.py
+++ b/templates/python/main.py
@@ -2,6 +2,7 @@ import argparse
 import logging
 import os
 import random
+import shutil
 import tarfile
 import tempfile
 
@@ -85,7 +86,7 @@ def main(scenarios):
 
     try:
         for scen in todo:
-            scen.run()
+            scen.run(_datadir)
     except Exception as e:
         logging.error(str(e), exc_info=True)
         if args.save_on_failure:
@@ -94,3 +95,7 @@ def main(scenarios):
             print(filename)
             save_directory(_datadir, filename)
         raise
+
+    shutil.rmtree(_datadir)
+    print("OK, all scenarios finished successfully")
+    logging.info("OK, all scenarios finished successfully")
diff --git a/templates/python/scenarios.py b/templates/python/scenarios.py
index b2738eb..71ed880 100644
--- a/templates/python/scenarios.py
+++ b/templates/python/scenarios.py
@@ -45,9 +45,10 @@ class Step:
 
 
 class Scenario:
-    def __init__(self):
+    def __init__(self, ctx):
         self._title = None
         self._steps = []
+        self._ctx = ctx
 
     def get_title(self):
         return self._title
@@ -58,16 +59,16 @@ class Scenario:
     def append_step(self, step):
         self._steps.append(step)
 
-    def run(self):
+    def run(self, datadir):
         print("scenario: {}".format(self._title))
         logging.info("Scenario: {}".format(self._title))
 
-        scendir = tempfile.mkdtemp(dir=_datadir)
+        scendir = tempfile.mkdtemp(dir=datadir)
         os.chdir(scendir)
         os.environ["HOME"] = scendir
 
-        ctx = Context()
         done = []
+        ctx = self._ctx
         try:
             for step in self._steps:
                 step.do(ctx)
diff --git a/templates/python/template.py b/templates/python/template.py
index 7b8ce45..f3f928b 100644
--- a/templates/python/template.py
+++ b/templates/python/template.py
@@ -12,10 +12,6 @@
 #############################################################################
 # Scaffolding for generated test program.
 
-# These imports are needed by the code in this template.
-import logging
-import shutil
-
 {% include "context.py" %}
 {% include "encoding.py" %}
 {% include "files.py" %}
@@ -28,12 +24,11 @@ import shutil
 # Test data files that were embedded in the source document. Base64
 # encoding is used to allow arbitrary data.
 
-_files = {}
 {% for file in files %}
 # {{ file.filename }}
 filename = decode_str('{{ file.filename | base64 }}')
 contents = decode_bytes('{{ file.contents | base64 }}')
-_files[filename] = contents
+store_file(filename, contents)
 {% endfor %}
 
 
@@ -46,7 +41,8 @@ _files[filename] = contents
 # Scenario: {{ scenario.title }}
 class Scenario_{{ loop.index }}():
     def __init__(self):
-        self._scenario = Scenario()
+        ctx = Context()
+        self._scenario = Scenario(ctx)
         self._scenario.set_title(decode_str('{{ scenario.title | base64 }}'))
         {% for step in scenario.steps %}
         # Step: {{ step.text }}
@@ -68,8 +64,8 @@ class Scenario_{{ loop.index }}():
     def get_title(self):
         return self._scenario.get_title()
 
-    def run(self):
-        self._scenario.run()
+    def run(self, datadir):
+        self._scenario.run(datadir)
 {% endfor %}
 
 _scenarios = { {% for scenario in scenarios %}
@@ -80,6 +76,3 @@ _scenarios = { {% for scenario in scenarios %}
 #############################################################################
 # Call main function and clean up.
 main(_scenarios)
-shutil.rmtree(_datadir)
-print('OK, all scenarios finished successfully')
-logging.info("OK, all scenarios finished successfully")