examples: Add swapout CLI argument to scalingtimeHEADmaster

Add some methods to the enum for nicer reporting when printing the help
documentation or when errors occur due to incorrect values on the
command line.

Signed-off-by: Thijs Paelman <thijs@ouroboros.rocks>

2 files changed, 15 insertions, 1 deletions

diff --git a/examples/scalingtime.py b/examples/scalingtime.py
index ff3c2ed..547622f 100755
--- a/examples/scalingtime.py
+++ b/examples/scalingtime.py
@@ -37,6 +37,10 @@ argparser.add_argument('--verbosity', metavar='VERBOSITY', type=str,
                        choices=['DEBUG', 'INFO', 'WARNING',
                                 'ERROR', 'CRITICAL'],
                        help='Logging verbosity level')
+argparser.add_argument('--swapout', type=SwapOutStrategy.from_string,
+                       default=PROMPT_SWAPOUT,
+                       choices=list(SwapOutStrategy),
+                       help='What action to perform on swap-out')
 argparser.add_argument('--prototype', type=str, required=True,
                        choices=['irati', 'ouroboros', 'rlite'],
                        help='The kind of prototype plugin to use to run'
@@ -239,7 +243,7 @@ exp = exp_class(testbed, nodes=nodes)
 
 if __name__ == "__main__":
     draw_experiment(exp)
-    with ExperimentManager(exp, swap_out_strategy=PROMPT_SWAPOUT):
+    with ExperimentManager(exp, swap_out_strategy=args.swapout):
         exp.swap_in()
         if not isinstance(testbed, docker.Testbed) \
                 and not isinstance(testbed, qemu.Testbed) \
diff --git a/rumba/utils.py b/rumba/utils.py
index 822deb3..aa06415 100644
--- a/rumba/utils.py
+++ b/rumba/utils.py
@@ -57,6 +57,16 @@ class SwapOutStrategy(enum.Enum):
     PAUSE = 2
     PROMPT = 3
 
+    def __str__(self):
+        return self.name
+
+    @staticmethod
+    def from_string(s):
+        try:
+            return SwapOutStrategy[s]
+        except KeyError:
+            raise ValueError()
+
 
 class SyslogsStrategy(enum.Enum):
     """