Merge pull request #136 from simvue-io/wk9874/get_alerts

Add `get_alerts` and `get_run_id_from_name` methods to Client

Author: kzscisoft

simvue/client.py

@@ -39,6 +39,24 @@ def __init__(self):
         self._headers = {"Authorization": f"Bearer {self._token}"}
         self._version = get_server_version()
 
+    def get_run_id_from_name(self, name):
+        """
+        Get run id for the specified run name
+        """
+        params = {'filters': json.dumps([f"name == {name}"])}
+
+        response = requests.get(f"{self._url}/api/runs", headers=self._headers, params=params)
+
+        if response.status_code == 200:
+            if 'data' in response.json():
+                if len(response.json()['data']) == 0:
+                    raise RuntimeError("Could not collect ID - no run found with this name.")
+                if len(response.json()['data']) > 1:
+                    raise RuntimeError("Could not collect ID - more than one run exists with this name.")
+                else:
+                    return response.json()['data'][0]['id']
+        raise RuntimeError(response.text)
+
     def get_run(self, run, system=False, tags=False, metadata=False):
         """
         Get a single run
@@ -468,3 +486,36 @@ def get_events(self, run, filter=None, start=0, num=0):
             return response.json()['data']
 
         raise Exception(response.text)
+    
+    def get_alerts(self, run, triggered_only = True, names_only = True):
+        """_summary_
+
+        Parameters
+        ----------
+        run : str
+            The ID of the run to find alerts for
+        critical_only : bool, optional
+            Whether to only return details about alerts which are currently critical, by default True
+        names_only: bool, optional
+            Whether to only return the names of the alerts (otherwise return the full details of the alerts), by default True
+        """
+        response = requests.get(f"{self._url}/api/runs/{run}", headers=self._headers)
+
+        if response.status_code == 404:
+            if 'detail' in response.json():
+                if response.json()['detail'] == 'run does not exist':
+                    raise Exception('Run does not exist')
+
+        elif response.status_code == 200:
+            if triggered_only:
+                if names_only:
+                    return [alert['alert']['name'] for alert in response.json()['alerts'] if alert['status']['current'] == 'critical']
+                else:
+                    return [alert for alert in response.json()['alerts'] if alert['status']['current'] == 'critical']
+            else:
+                if names_only:
+                    return [alert['alert']['name'] for alert in response.json()['alerts']]
+                else:
+                    return response.json()['alerts']
+
+        raise Exception(response.text)

tests/functional/test_get_alerts.py

@@ -0,0 +1,60 @@
+import time
+import unittest
+import uuid
+from simvue import Run, Client
+from simvue.sender import sender
+
+import common
+
+class TestGetAlerts(unittest.TestCase):
+    def test_get_alerts(self):
+        """
+        Create a run & two alerts, make one alert be triggered
+        If triggered_only is True, check that only the triggered alert is returned
+        IF triggered_only is False, check both alerts are returned
+        If names_only is False, check that full dictionary of information is returned
+        """
+        run = Run()
+        client = Client()
+        folder = '/test-%s' % str(uuid.uuid4())
+        name = common.RUNNAME3
+        run.init(name, folder=folder)
+        run.add_alert(
+            name='value_below_1',
+            source='metrics',
+            frequency=1,
+            rule='is below',
+            threshold=1,
+            metric='test_metric',
+            window=2
+        )
+        run.add_alert(
+            name='value_above_1',
+            source='metrics',
+            frequency=1,
+            rule='is above',
+            threshold=1,
+            metric='test_metric',
+            window=2
+        )
+
+        run.log_metrics({'test_metric': 5})
+        time.sleep(180)
+
+        run_id = client.get_run_id_from_name(name)
+
+        triggered_alerts_names = client.get_alerts(run_id)
+        self.assertListEqual(triggered_alerts_names, ['value_above_1'])
+
+        triggered_alerts_full = client.get_alerts(run_id, names_only=False)
+        self.assertIsInstance(triggered_alerts_full[0], dict)
+        self.assertEqual(triggered_alerts_full[0]["alert"]["name"], "value_above_1")
+        self.assertEqual(triggered_alerts_full[0]["status"]["current"], "critical")
+
+        all_alerts_names = client.get_alerts(run_id, triggered_only=False)
+        self.assertListEqual(all_alerts_names, ['value_above_1', 'value_below_1'])
+
+        run.close()
+
+if __name__ == '__main__':
+    unittest.main()