refactor(templates/python): split out helper code into modules

Also add unit tests.

10 files changed, 305 insertions, 231 deletions

diff --git a/templates/python/asserts.py b/templates/python/asserts.py
new file mode 100644
index 0000000..c898454
--- /dev/null
+++ b/templates/python/asserts.py
@@ -0,0 +1,23 @@
+# Check two values for equality and give error if they are not equal
+def assert_eq(a, b):
+    assert a == b, "expected %r == %r" % (a, b)
+
+
+# Check two values for inequality and give error if they are equal
+def assert_ne(a, b):
+    assert a != b, "expected %r != %r" % (a, b)
+
+
+# Check that two dict values are equal.
+def assert_dict_eq(a, b):
+    assert isinstance(a, dict)
+    assert isinstance(b, dict)
+    for key in a:
+        assert key in b, f"exected {key} in both dicts"
+        av = a[key]
+        bv = b[key]
+        assert_eq(type(av), type(bv))
+        if isinstance(av, list):
+            assert_eq(list(sorted(av)), list(sorted(bv)))
+    for key in b:
+        assert key in a, f"exected {key} in both dicts"
diff --git a/templates/python/context.py b/templates/python/context.py
new file mode 100644
index 0000000..8d9894f
--- /dev/null
+++ b/templates/python/context.py
@@ -0,0 +1,20 @@
+import logging
+
+
+# Store context between steps.
+class Context:
+    def __init__(self):
+        self._vars = {}
+
+    def as_dict(self):
+        return dict(self._vars)
+
+    def get(self, key, default=None):
+        return self._vars.get(key, default)
+
+    def __getitem__(self, key):
+        return self._vars[key]
+
+    def __setitem__(self, key, value):
+        logging.info("Context: {}={!r}".format(key, value))
+        self._vars[key] = value
diff --git a/templates/python/context_tests.py b/templates/python/context_tests.py
new file mode 100644
index 0000000..fc99af5
--- /dev/null
+++ b/templates/python/context_tests.py
@@ -0,0 +1,30 @@
+import unittest
+
+from context import Context
+
+
+class ContextTests(unittest.TestCase):
+    def test_converts_to_empty_dict_initially(self):
+        ctx = Context()
+        self.assertEqual(ctx.as_dict(), {})
+
+    def test_set_item(self):
+        ctx = Context()
+        ctx["foo"] = "bar"
+        self.assertEqual(ctx["foo"], "bar")
+
+    def test_get_returns_default_if_item_does_not_exist(self):
+        ctx = Context()
+        self.assertEqual(ctx.get("foo"), None)
+
+    def test_get_returns_specified_default_if_item_does_not_exist(self):
+        ctx = Context()
+        self.assertEqual(ctx.get("foo", "bar"), "bar")
+
+    def test_get_returns_value_if_item_exists(self):
+        ctx = Context()
+        ctx["foo"] = "bar"
+        self.assertEqual(ctx.get("foo", "yo"), "bar")
+
+
+unittest.main()
diff --git a/templates/python/encoding.py b/templates/python/encoding.py
new file mode 100644
index 0000000..1efb95e
--- /dev/null
+++ b/templates/python/encoding.py
@@ -0,0 +1,12 @@
+# Decode a base64 encoded string. Result is binary or unicode string.
+
+
+import base64
+
+
+def decode_bytes(s):
+    return base64.b64decode(s)
+
+
+def decode_str(s):
+    return base64.b64decode(s).decode()
diff --git a/templates/python/encoding_tests.py b/templates/python/encoding_tests.py
new file mode 100644
index 0000000..4167aa4
--- /dev/null
+++ b/templates/python/encoding_tests.py
@@ -0,0 +1,19 @@
+import base64
+import unittest
+
+import encoding
+
+
+class EncodingTests(unittest.TestCase):
+    def test_str_roundtrip(self):
+        original = "foo\nbar\0"
+        encoded = base64.b64encode(original.encode())
+        self.assertEqual(encoding.decode_str(encoded), original)
+
+    def test_bytes_roundtrip(self):
+        original = b"foo\nbar\0"
+        encoded = base64.b64encode(original)
+        self.assertEqual(encoding.decode_bytes(encoded), original)
+
+
+unittest.main()
diff --git a/templates/python/files.py b/templates/python/files.py
new file mode 100644
index 0000000..9adbd30
--- /dev/null
+++ b/templates/python/files.py
@@ -0,0 +1,3 @@
+# Retrieve an embedded test data file using filename.
+def get_file(filename):
+    return _files[filename]
diff --git a/templates/python/main.py b/templates/python/main.py
new file mode 100644
index 0000000..3a39164
--- /dev/null
+++ b/templates/python/main.py
@@ -0,0 +1,96 @@
+import argparse
+import logging
+import os
+import random
+import tarfile
+import tempfile
+
+
+# Remember where we started from. The step functions may need to refer
+# to files there.
+srcdir = os.getcwd()
+print("srcdir", srcdir)
+
+# Create a new temporary directory and chdir there. This allows step
+# functions to create new files in the current working directory
+# without having to be so careful.
+_datadir = tempfile.mkdtemp()
+print("datadir", _datadir)
+os.chdir(_datadir)
+
+
+def parse_command_line():
+    p = argparse.ArgumentParser()
+    p.add_argument("--log")
+    p.add_argument("--save-on-failure")
+    p.add_argument("patterns", nargs="*")
+    return p.parse_args()
+
+
+def setup_minimal_environment():
+    minimal = {
+        "PATH": "/bin:/usr/bin",
+        "SHELL": "/bin/sh",
+        "HOME": "/",  # will be set to datadir for each scenario
+    }
+
+    os.environ.clear()
+    os.environ.update(minimal)
+
+
+def setup_logging(args):
+    if args.log:
+        fmt = "%(asctime)s %(levelname)s %(message)s"
+        datefmt = "%Y-%m-%d %H:%M:%S"
+        formatter = logging.Formatter(fmt, datefmt)
+
+        filename = os.path.abspath(os.path.join(srcdir, args.log))
+        handler = logging.FileHandler(filename)
+        handler.setFormatter(formatter)
+    else:
+        handler = logging.NullHandler()
+
+    logger = logging.getLogger()
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+
+
+def save_directory(dirname, tarname):
+    print("tarname", tarname)
+    logging.info("Saving {} to {}".format(dirname, tarname))
+    tar = tarfile.open(tarname, "w")
+    tar.add(dirname, arcname="datadir")
+    tar.close()
+
+
+def main(scenarios):
+    args = parse_command_line()
+    setup_minimal_environment()
+    setup_logging(args)
+    logging.info("Test program starts")
+
+    logging.info("patterns: {}".format(args.patterns))
+    if len(args.patterns) == 0:
+        logging.info("Executing all scenarios")
+        todo = list(scenarios)
+        random.shuffle(todo)
+    else:
+        logging.info("Executing requested scenarios only: {}".format(args.patterns))
+        patterns = [arg.lower() for arg in args.patterns]
+        todo = [
+            scen
+            for scen in scenarios
+            if any(pattern in scen.get_title().lower() for pattern in patterns)
+        ]
+
+    try:
+        for scen in todo:
+            scen.run()
+    except Exception as e:
+        logging.error(str(e), exc_info=True)
+        if args.save_on_failure:
+            print(args.save_on_failure)
+            filename = os.path.abspath(os.path.join(srcdir, args.save_on_failure))
+            print(filename)
+            save_directory(_datadir, filename)
+        raise
diff --git a/templates/python/scenarios.py b/templates/python/scenarios.py
new file mode 100644
index 0000000..b2738eb
--- /dev/null
+++ b/templates/python/scenarios.py
@@ -0,0 +1,81 @@
+import logging
+import os
+import tempfile
+
+
+#############################################################################
+# Code to implement the scenarios.
+
+
+class Step:
+    def __init__(self):
+        self._kind = None
+        self._text = None
+        self._args = {}
+        self._function = None
+        self._cleanup = 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 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:
+    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)
+        os.environ["HOME"] = scendir
+
+        ctx = Context()
+        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)
diff --git a/templates/python/template.py b/templates/python/template.py
index af5b36e..7b8ce45 100644
--- a/templates/python/template.py
+++ b/templates/python/template.py
@@ -10,47 +10,24 @@
 
 
 #############################################################################
-# Helper code generated by Subplot.
+# Scaffolding for generated test program.
 
-import argparse
-import base64
+# These imports are needed by the code in this template.
 import logging
-import os
-import random
 import shutil
-import sys
-import tarfile
-import tempfile
 
-# Store context between steps.
-class Context:
+{% include "context.py" %}
+{% include "encoding.py" %}
+{% include "files.py" %}
+{% include "asserts.py" %}
+{% include "scenarios.py" %}
+{% include "main.py" %}
 
-    def __init__(self):
-        self._vars = {}
-
-    def as_dict(self):
-        return dict(self._vars)
-
-    def get(self, key, default=None):
-        return self._vars.get(key, default)
-
-    def __getitem__(self, key):
-        return self._vars[key]
-
-    def __setitem__(self, key, value):
-        logging.info('Context: {}={!r}'.format(key, value))
-        self._vars[key] = value
-
-# Decode a base64 encoded string. Result is binary or unicode string.
-
-def decode_bytes(s):
-    return base64.b64decode(s)
-
-def decode_str(s):
-    return base64.b64decode(s).decode()
 
+#############################################################################
 # 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 }}
@@ -60,126 +37,12 @@ _files[filename] = contents
 {% endfor %}
 
 
-# Retrieve an embedded test data file using filename.
-def get_file(filename):
-    return _files[filename]
-
-# Check two values for equality and give error if they are not equal
-def assert_eq(a, b):
-    assert a == b, 'expected %r == %r' % (a, b)
-
-# Check two values for inequality and give error if they are equal
-def assert_ne(a, b):
-    assert a != b, 'expected %r != %r' % (a, b)
-
-# Check that two dict values are equal.
-def assert_dict_eq(a, b):
-    assert isinstance(a, dict)
-    assert isinstance(b, dict)
-    for key in a:
-        assert key in b, f"exected {key} in both dicts"
-        av = a[key]
-        bv = b[key]
-        assert_eq(type(av), type(bv))
-        if isinstance(av, list):
-            assert_eq(list(sorted(av)), list(sorted(bv)))
-    for key in b:
-        assert key in a, f"exected {key} in both dicts"
-
-# Remember where we started from. The step functions may need to refer
-# to files there.
-srcdir = os.getcwd()
-print('srcdir', srcdir)
-
-# Create a new temporary directory and chdir there. This allows step
-# functions to create new files in the current working directory
-# without having to be so careful.
-_datadir = tempfile.mkdtemp()
-print('datadir', _datadir)
-os.chdir(_datadir)
 
 #############################################################################
-# Code to implement the scenarios.
-
-
-class Step:
-
-    def __init__(self):
-        self._kind = None
-        self._text = None
-        self._args = {}
-        self._function = None
-        self._cleanup = 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 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:
-
-    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)
-        os.environ['HOME'] = scendir
-
-        ctx = Context()
-        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)
-
+# Classes for individual scenarios.
 
 {% for scenario in scenarios %}
-######################################
+#----------------------------------------------------------------------------
 # Scenario: {{ scenario.title }}
 class Scenario_{{ loop.index }}():
     def __init__(self):
@@ -214,88 +77,9 @@ _scenarios = { {% for scenario in scenarios %}
 }
 
 
-def parse_command_line():
-    p = argparse.ArgumentParser()
-    p.add_argument("--log")
-    p.add_argument("--save-on-failure")
-    p.add_argument("patterns", nargs="*")
-    return p.parse_args()
-
-
-def setup_minimal_environment():
-    minimal = {
-        'PATH': '/bin:/usr/bin',
-        'SHELL': '/bin/sh',
-        'HOME': '/',  # will be set to datadir for each scenario
-    }
-
-    os.environ.clear()
-    os.environ.update(minimal)
-
-
-def setup_logging(args):
-    if args.log:
-        fmt = "%(asctime)s %(levelname)s %(message)s"
-        datefmt = "%Y-%m-%d %H:%M:%S"
-        formatter = logging.Formatter(fmt, datefmt)
-
-        filename = os.path.abspath(os.path.join(srcdir, args.log))
-        handler = logging.FileHandler(filename)
-        handler.setFormatter(formatter)
-    else:
-        handler = logging.NullHandler()
-
-    logger = logging.getLogger()
-    logger.addHandler(handler)
-    logger.setLevel(logging.DEBUG)
-
-
-def save_directory(dirname, tarname):
-    print('tarname', tarname)
-    logging.info("Saving {} to {}".format(dirname, tarname))
-    tar = tarfile.open(tarname, "w")
-    tar.add(dirname, arcname="datadir")
-    tar.close()
-
-
-def main():
-    args = parse_command_line()
-    setup_minimal_environment()
-    setup_logging(args)
-    logging.info("Test program starts")
-
-    logging.info("patterns: {}".format(args.patterns))
-    if len(args.patterns) == 0:
-        logging.info("Executing all scenarios")
-        todo = list(_scenarios)
-        random.shuffle(todo)
-    else:
-        logging.info("Executing requested scenarios only: {}".format(args.patterns))
-        patterns = [arg.lower() for arg in args.patterns]
-        todo = [
-            scen
-            for scen in _scenarios
-            if any(pattern in scen.get_title().lower() for pattern in patterns)
-        ]
-
-    try:
-        for scen in todo:
-            scen.run()
-    except Exception as e:
-        logging.error(str(e), exc_info=True)
-        if args.save_on_failure:
-            print(args.save_on_failure)
-            filename = os.path.abspath(os.path.join(srcdir, args.save_on_failure))
-            print(filename)
-            save_directory(_datadir, filename)
-        raise
-
-
-main()
-
 #############################################################################
-# Clean up temporary directory and report success.
-
+# Call main function and clean up.
+main(_scenarios)
 shutil.rmtree(_datadir)
 print('OK, all scenarios finished successfully')
 logging.info("OK, all scenarios finished successfully")
diff --git a/templates/python/template.yaml b/templates/python/template.yaml
index c3ded5a..afb6522 100644
--- a/templates/python/template.yaml
+++ b/templates/python/template.yaml
@@ -1,3 +1,9 @@
 template: template.py
-helpers: []
+helpers:
+  - context.py
+  - encoding.py
+  - files.py
+  - asserts.py
+  - scenarios.py
+  - main.py
 run: python3