add scenario

3 files changed, 18 insertions, 6 deletions

diff --git a/templates/python/main.py b/templates/python/main.py
index 05e1715..87e2782 100644
--- a/templates/python/main.py
+++ b/templates/python/main.py
@@ -23,6 +23,7 @@ os.chdir(_datadir)
 def parse_command_line():
     p = argparse.ArgumentParser()
     p.add_argument("--log")
+    p.add_argument("--env", action="append", default=[])
     p.add_argument("--save-on-failure")
     p.add_argument("patterns", nargs="*")
     return p.parse_args()
@@ -72,9 +73,16 @@ def main(scenarios):
             if any(pattern in scen.get_title().lower() for pattern in patterns)
         ]
 
+    extra_env = {}
+    for env in args.env:
+        (name, value) = env.split("=", 1)
+        extra_env[name] = value
+    logging.debug(f"args.env: {args.env}")
+    logging.debug(f"env vars from command line; {extra_env}")
+
     try:
         for scen in todo:
-            scen.run(_datadir)
+            scen.run(_datadir, extra_env)
     except Exception as e:
         logging.error(str(e), exc_info=True)
         if args.save_on_failure:
diff --git a/templates/python/scenarios.py b/templates/python/scenarios.py
index eb5302e..e2703df 100644
--- a/templates/python/scenarios.py
+++ b/templates/python/scenarios.py
@@ -59,13 +59,14 @@ class Scenario:
     def append_step(self, step):
         self._steps.append(step)
 
-    def run(self, datadir):
+    def run(self, datadir, extra_env):
         print("scenario: {}".format(self._title))
         logging.info("Scenario: {}".format(self._title))
+        logging.info("extra environment variables: {}".format(extra_env))
 
         scendir = tempfile.mkdtemp(dir=datadir)
         os.chdir(scendir)
-        self._set_environment_variables_to(scendir)
+        self._set_environment_variables_to(scendir, extra_env)
 
         done = []
         ctx = self._ctx
@@ -81,7 +82,7 @@ class Scenario:
         for step in reversed(done):
             step.cleanup(ctx)
 
-    def _set_environment_variables_to(self, scendir):
+    def _set_environment_variables_to(self, scendir, extra_env):
         minimal = {
             "PATH": "/bin:/usr/bin",
             "SHELL": "/bin/sh",
@@ -91,3 +92,6 @@ class Scenario:
 
         os.environ.clear()
         os.environ.update(minimal)
+        os.environ.update(extra_env)
+        logging.debug(f"extra_env: {dict(extra_env)!r}")
+        logging.debug(f"os.environ: {dict(os.environ)!r}")
diff --git a/templates/python/template.py b/templates/python/template.py
index f3f928b..aa97cf0 100644
--- a/templates/python/template.py
+++ b/templates/python/template.py
@@ -64,8 +64,8 @@ class Scenario_{{ loop.index }}():
     def get_title(self):
         return self._scenario.get_title()
 
-    def run(self, datadir):
-        self._scenario.run(datadir)
+    def run(self, datadir, extra_env):
+        self._scenario.run(datadir, extra_env)
 {% endfor %}
 
 _scenarios = { {% for scenario in scenarios %}