storyboard: add callback parameter to run_client(_of) method

1 files changed, 29 insertions, 8 deletions

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