sb + our: add server_decorator functionality and our decorator

1 files changed, 9 insertions, 1 deletions

diff --git a/rumba/elements/experimentation.py b/rumba/elements/experimentation.py
index 8f3f40a..0f734f8 100644
--- a/rumba/elements/experimentation.py
+++ b/rumba/elements/experimentation.py
@@ -117,7 +117,8 @@ class Experiment(object):
                  log_dir=None,
                  prototype_logs=None,
                  enrollment_strategy='minimal',
-                 dt_strategy='full-mesh'):
+                 dt_strategy='full-mesh',
+                 server_decorator=None):
         """
         :param testbed: The testbed of the experiment.
         :param nodes: The list of nodes in the experiment.
@@ -127,10 +128,17 @@ class Experiment(object):
         :param prototype_logs: Where the prototype logs its output.
         :param enrollment_strategy: Can be 'full-mesh', 'minimal' or 'manual'.
         :param dt_strategy: For data flows, 'full-mesh', 'minimal' or 'manual'.
+        :param server_decorator: a decorator function which will be applied to
+                                 storyboard.server instances when using
+                                 this prototype
         """
         if nodes is None:
             nodes = list()
         self.nodes = nodes
+        if server_decorator is None:
+            def server_decorator(server):
+                return server
+        self.server_decorator = server_decorator
         self.git_repo = git_repo
         self.git_branch = git_branch
         self.testbed = testbed