elements: Experiment: add missing enrollment_strategy argument

This is necessary to enable scripts to use full-mesh or manual
enrollment strategies.
Fixes #54.

4 files changed, 21 insertions, 8 deletions

diff --git a/rumba/elements/experimentation.py b/rumba/elements/experimentation.py
index da81c6c..990f285 100644
--- a/rumba/elements/experimentation.py
+++ b/rumba/elements/experimentation.py
@@ -115,7 +115,8 @@ class Experiment(object):
                  git_repo=None,
                  git_branch=None,
                  log_dir=None,
-                 prototype_logs=None):
+                 prototype_logs=None,
+                 enrollment_strategy='minimal'):
         """
         :param testbed: The testbed of the experiment.
         :param nodes: The list of nodes in the experiment.
@@ -123,6 +124,7 @@ class Experiment(object):
         :param git_branch: The git branch of the repository.
         :param log_dir: Where to log output of the experiment.
         :param prototype_logs: Where the prototype logs its output.
+        :param enrollment_strategy: Can be 'full-mesh', 'minimal' or 'manual'.
         """
         if nodes is None:
             nodes = list()
@@ -132,7 +134,7 @@ class Experiment(object):
         self.testbed = testbed
         # the strategy employed for completing the enrollment phase in
         # the different DIFs
-        self.enrollment_strategy = 'minimal'  # 'full-mesh', 'manual'
+        self.enrollment_strategy = enrollment_strategy  # 'full-mesh', 'manual'
         # the strategy employed for setting up the data transfer
         # networks in the DIFs after enrollment
         self.dt_strategy = 'full-mesh'  # 'minimal', 'manual'
@@ -141,6 +143,10 @@ class Experiment(object):
         self.dt_flows = []  # a list of per-DIF lists of data transfer flows
         self.mgmt_flows = []  # a list of per-DIF lists of management flows
 
+        if self.enrollment_strategy not in ['full-mesh', 'minimal', 'manual']:
+            raise Exception('Unknown enrollment strategy "%s"'
+                                % enrollment_strategy)
+
         # Determine log directory
         if log_dir is None:
             # If it is None, use /tmp/rumba/{project}
diff --git a/rumba/prototypes/irati.py b/rumba/prototypes/irati.py
index 1967278..ec3e53d 100644
--- a/rumba/prototypes/irati.py
+++ b/rumba/prototypes/irati.py
@@ -67,7 +67,8 @@ class Experiment(mod.Experiment):
                  git_branch='arcfire',
                  installpath=None,
                  varpath=None,
-                 app_mappings=None):
+                 app_mappings=None,
+                 enrollment_strategy='minimal'):
         """
         Initializes the experiment class.
 
@@ -81,13 +82,15 @@ class Experiment(mod.Experiment):
                              all application which will register to
                              any given dif.
         :type app_mappings: `List[(str, str)]`
+        :param enrollment_strategy: Can be 'full-mesh', 'minimal' or 'manual'.
         """
         mod.Experiment.__init__(self,
                                 testbed,
                                 nodes,
                                 git_repo,
                                 git_branch,
-                                prototype_logs=['/tmp/ipcmstart.log'])
+                                prototype_logs=['/tmp/ipcmstart.log'],
+                                enrollment_strategy=enrollment_strategy)
         if installpath is None:
             installpath = '/usr'
         if varpath is None:
diff --git a/rumba/prototypes/ouroboros.py b/rumba/prototypes/ouroboros.py
index e72bcf1..f415160 100644
--- a/rumba/prototypes/ouroboros.py
+++ b/rumba/prototypes/ouroboros.py
@@ -44,7 +44,7 @@ class Experiment(mod.Experiment):
     """
     def __init__(self, testbed, nodes=None,
                  git_repo='git://ouroboros.ilabt.imec.be/ouroboros',
-                 git_branch='master'):
+                 git_branch='master', enrollment_strategy='minimal'):
         """
         Initializes the experiment class.
 
@@ -52,8 +52,10 @@ class Experiment(mod.Experiment):
         :param nodes: The list of nodes.
         :param git_repo: The git repository to use for installation.
         :param git_branch: The branch of the git repository to use.
+        :param enrollment_strategy: Can be 'full-mesh', 'minimal' or 'manual'.
         """
-        mod.Experiment.__init__(self, testbed, nodes, git_repo, git_branch)
+        mod.Experiment.__init__(self, testbed, nodes, git_repo, git_branch,
+                                enrollment_strategy=enrollment_strategy)
         self.r_ipcps = dict()
 
         self.set_startup_command("irmd")
diff --git a/rumba/prototypes/rlite.py b/rumba/prototypes/rlite.py
index 32b7df7..1595336 100644
--- a/rumba/prototypes/rlite.py
+++ b/rumba/prototypes/rlite.py
@@ -42,7 +42,7 @@ class Experiment(mod.Experiment):
 
     def __init__(self, testbed, nodes=None,
                  git_repo='https://github.com/vmaffione/rlite',
-                 git_branch='master'):
+                 git_branch='master', enrollment_strategy='minimal'):
         """
         Initializes the experiment class.
 
@@ -50,9 +50,11 @@ class Experiment(mod.Experiment):
         :param nodes: The list of nodes.
         :param git_repo: The git repository to use for installation.
         :param git_branch: The branch of the git repository to use.
+        :param enrollment_strategy: Can be 'full-mesh', 'minimal' or 'manual'.
         """
         mod.Experiment.__init__(self, testbed, nodes, git_repo, git_branch,
-                                prototype_logs=['/tmp/uipcp.log'])
+                                prototype_logs=['/tmp/uipcp.log'],
+                                enrollment_strategy=enrollment_strategy)
 
     @staticmethod
     def make_executor(node, packages, testbed):