Change: Python template to run only chosen scenarios

1 files changed, 21 insertions, 6 deletions

diff --git a/templates/python/template.py b/templates/python/template.py
index a0f4ec5..87ced78 100644
--- a/templates/python/template.py
+++ b/templates/python/template.py
@@ -12,6 +12,7 @@ import logging
 import os
 import random
 import shutil
+import sys
 import tempfile
 
 # Store context between steps.
@@ -102,13 +103,27 @@ def scenario_{{ loop.index }}():
     {% endfor %}
 {% endfor %}
 
+_scenarios = {
+{% for scenario in scenarios %}
+    '{{ scenario.title }}': scenario_{{ loop.index }},
+{% endfor %}
+}
+
+
 def main():
-    scenarios = [{% for scenario in scenarios %}
-        scenario_{{ loop.index }},{% endfor %}
-    ]
-    random.shuffle(scenarios)
-    for s in scenarios:
-        s()
+    if len(sys.argv) == 1:
+        funcs = list(_scenarios.values())
+        random.shuffle(funcs)
+    else:
+        patterns = [arg.lower() for arg in sys.argv[1:]]
+        funcs = [
+            func
+            for title, func in _scenarios.items()
+            if any(pattern in title.lower() for pattern in patterns)
+        ]
+
+    for func in funcs:
+        func()
 
 main()