Added run and wait method for scenarios

Author: nicolasservel

dataikuapi/dss/scenario.py

@@ -1,4 +1,7 @@
 from datetime import datetime
+import time
+from dataikuapi.utils import DataikuException
+
 
 class DSSScenario(object):
     """
@@ -26,6 +29,50 @@ def run(self, params={}):
             "POST", "/projects/%s/scenarios/%s/run" % (self.project_key, self.id), body=params)
         return DSSTriggerFire(self, trigger_fire)
 
+    def get_trigger(self, trigger_id, run_id):
+        """
+        Requests a trigger of the run of a scenario
+
+        Args:
+            trigger_id:  Id of trigger
+            run_id: Id of associated run
+
+        Returns:
+            A :class:`dataikuapi.dss.admin.DSSTriggerFire` trigger handle
+        """
+        trigger_fire = self.client._perform_json(
+            "GET", "/projects/%s/scenarios/%s/get-trigger/%s" % (self.project_key, self.id, trigger_id), params={
+                'runId' : run_id
+            })
+        return DSSTriggerFire(self, trigger_fire)
+
+    def run_and_wait(self, params={}):
+        """
+        Requests a run of the scenario, which will start after a few seconds. Wait the end of the run to complete.
+
+        Args:
+            params: additional parameters that will be passed to the scenario through trigger params
+
+        Returns:
+            A :class:`dataikuapi.dss.admin.DSSScenarioRun` run handle
+        """
+        trigger_fire = self.run(params)
+        scenario_run = None
+        refresh_trigger_counter = 0
+        while scenario_run is None:
+            refresh_trigger_counter += 1
+            if refresh_trigger_counter == 10:
+                refresh_trigger_counter = 0
+                trigger_fire = self.get_trigger(trigger_fire.trigger_id, trigger_fire.run_id)
+            if trigger_fire.is_cancelled():
+                raise DataikuException("Scenario run has been cancelled")
+            scenario_run = trigger_fire.get_scenario_run()
+            time.sleep(5)
+        while not scenario_run.run.get('result', False):
+            scenario_run = trigger_fire.get_scenario_run()
+            time.sleep(60)
+        return scenario_run
+
     def get_last_runs(self, limit=10, only_finished_runs=False):
         """
         Get the list of the last runs of the scenario.
@@ -191,3 +238,6 @@ def get_scenario_run(self):
             return None
         else:
             return DSSScenarioRun(self.client, run['scenarioRun'])
+
+    def is_cancelled(self):
+        return self.trigger_fire["cancelled"]