storyboard: add callback parameter to run_client(_of) method

2 files changed, 40 insertions, 8 deletions

diff --git a/examples/script-example.py b/examples/script-example.py
index 5e14823..6cfaf42 100755
--- a/examples/script-example.py
+++ b/examples/script-example.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 # An example script leveraging the storyboard scripting functionality
+from functools import partial
 
 from rumba.storyboard import *
 from rumba.model import *
@@ -83,6 +84,11 @@ server_c = Server(
     s_id='server_c'
 )
 
+typewriter = Client(
+    "touch test_file",
+    shutdown=""
+)
+
 
 if __name__ == '__main__':
     story = StoryBoard(30)
@@ -98,6 +104,11 @@ if __name__ == '__main__':
         'example-script.rsb'
     )
     story.parse_script_file(script_file)
+    cb = partial(story.run_command, node_a, "ls -l test_file")
+    # This will return a file, because it will run after
+    # Triggering_event
+    action = partial(story.run_client, typewriter, 0.2, node=node_a, callback=cb)
+    story.add_event(Event(action, ev_time=12, ev_id='Triggering_event'))
     log.flush_log()
     with ExperimentManager(exp):
         exp.swap_in()
diff --git a/rumba/storyboard.py b/rumba/storyboard.py
index b98cc6f..e38a102 100644
--- a/rumba/storyboard.py
+++ b/rumba/storyboard.py
@@ -475,7 +475,12 @@ class StoryBoard(SBEntity):
         """
         self._script.del_event(event)
 
-    def run_client_of(self, server, duration=None, node=None, proc_id=None):
+    def run_client_of(self,
+                      server,
+                      duration=None,
+                      node=None,
+                      proc_id=None,
+                      callback=None):
         """
         Runs a random client of the specified server
         with the specified parameters.
@@ -488,15 +493,22 @@ class StoryBoard(SBEntity):
         @param duration: the duration of the client process
         @param node: (Node or str) the node on which the client should be run
         @param proc_id: the entity ID to use for the process
+        @param callback: callable or list thereof to be run
+                         after client termination
         """
         if isinstance(server, str):
             server = self.server_apps[server]
         if duration is None:
             duration = server.get_duration()
         client = random.choice(server.clients)
-        self.run_client(client, duration, node, proc_id)
-
-    def run_client(self, client, duration, node=None, proc_id=None):
+        self.run_client(client, duration, node, proc_id, callback)
+
+    def run_client(self,
+                   client,
+                   duration,
+                   node=None,
+                   proc_id=None,
+                   callback=None):
         """
         Runs the specified client app with the specified parameters.
 
@@ -507,6 +519,8 @@ class StoryBoard(SBEntity):
         @param duration: the duration of the client process
         @param node: (Node or str) the node on which the client should be run
         @param proc_id: the entity ID to use for the process
+        @param callback: callable or list thereof to be run
+                         after client termination
         """
         if isinstance(client, str):
             client = self.client_apps[client]
@@ -517,11 +531,19 @@ class StoryBoard(SBEntity):
             node = random.choice(client.nodes)
         elif isinstance(node, str):
             node = self.node_map[node]
+        if callback is None:
+            callback = []
+        elif not hasattr(callback, '__len__'):
+            callback = [callback]
         process = client.process(duration, node, proc_id)
         self.process_dict[process.id] = process
         process.run()
         action = functools.partial(self.kill_process, process.id)
-        self._script.add_event(Event(action, ev_time=(self.cur_time + duration)))
+        term_ev = Event(action, ev_time=(self.cur_time + duration))
+        self.add_event(term_ev)
+        for cb in callback:
+            cb_event = Event(action=cb, trigger=term_ev)
+            self.add_event(cb_event)
 
     def start_client_of(self, server, duration=None, node=None, proc_id=None):
         """
@@ -813,10 +835,10 @@ class Event(object):
         @param ev_time: (float) seconds to wait before running the event
         @param trigger: (Event) Event which must complete before this runs
         """
+        self.id = ev_id if ev_id is not None else self.generate_id()
         if ev_time is None and trigger is None:
             raise ValueError('No condition specified for event %s.' %
-                             ev_id)
-        self.id = ev_id if ev_id is not None else self.generate_id()
+                             self.id)
         self.action = action
         self.time = ev_time
         self._trigger = trigger
@@ -1116,7 +1138,6 @@ class _Script(object):
         for event in ev_list:
             buffer.write(repr(event) + '\n')
 
-
     def _parse_conditions(self, *conditions):
         """
         Parses condition strings and returns the conditions