[Workflows] Fix run notifications (mlrun#5783)

Author: GiladShapira94

mlrun/utils/notifications/notification/webhook.py

@@ -69,6 +69,29 @@ async def push(
             request_body["custom_html"] = custom_html
 
         if override_body:
+            list_edit_runs = []
+            for run in runs:
+                if hasattr(run, "to_dict"):
+                    run = run.to_dict()
+                if isinstance(run, dict):
+                    parsed_run = {
+                        "project": run["metadata"]["project"],
+                        "name": run["metadata"]["name"],
+                        "host": run["metadata"]["labels"]["host"],
+                        "status": {"state": run["status"]["state"]},
+                    }
+                    if run["status"].get("error", None):
+                        parsed_run["status"]["error"] = run["status"]["error"]
+                    elif run["status"].get("results", None):
+                        parsed_run["status"]["results"] = run["status"]["results"]
+                    list_edit_runs.append(parsed_run)
+            runs_value = str(list_edit_runs)
+            if isinstance(override_body, dict):
+                for key, value in override_body.items():
+                    if "{{ runs }}" in value:
+                        override_body[key] = value.replace("{{ runs }}", runs_value)
+                    elif "{{runs}}" in value:
+                        override_body[key] = value.replace("{{runs}}", runs_value)
             request_body = override_body
 
         # Specify the `verify_ssl` parameter value only for HTTPS urls.

mlrun/utils/notifications/notification_pusher.py

@@ -397,7 +397,7 @@ def _add_run_step(_step: mlrun_pipelines.models.PipelineStep):
             try:
                 _run = db.list_runs(
                     project=run.metadata.project,
-                    labels=f"mlrun_constants.MLRunInternalLabels.runner_pod={_step.node_name}",
+                    labels=f"{mlrun_constants.MLRunInternalLabels.runner_pod}={_step.node_name}",
                 )[0]
             except IndexError:
                 _run = {

tests/utils/test_notifications.py

@@ -29,6 +29,7 @@
 import server.api.api.utils
 import server.api.constants
 import server.api.crud
+from mlrun.utils.notifications.notification.webhook import WebhookNotification
 
 
 @pytest.mark.parametrize(
@@ -196,6 +197,46 @@ def test_condition_evaluation_timeout():
     assert notification_pusher._should_notify(run, notification)
 
 
+@pytest.mark.parametrize(
+    "override_body",
+    [({"message": "runs: {{runs}}"}), ({"message": "runs: {{ runs }}"})],
+)
+async def test_webhook_overide_body_job_succeed(monkeypatch, override_body):
+    requests_mock = _mock_async_response(monkeypatch, "post", {"id": "response-id"})
+    runs = _generate_run_result(state="completed", results={"return": 1})
+    await WebhookNotification(
+        params={"override_body": override_body, "url": "http://test.com"}
+    ).push("test-message", "info", [runs])
+    expected_body = {
+        "message": "runs: [{'project': 'test-remote-workflow', 'name': 'func-func', 'host': 'func-func-8lvl8', "
+        "'status': {'state': 'completed', 'results': {'return': 1}}}]"
+    }
+    requests_mock.assert_called_once_with(
+        "http://test.com", headers={}, json=expected_body, ssl=None
+    )
+
+
+@pytest.mark.parametrize(
+    "override_body",
+    [({"message": "runs: {{runs}}"}), ({"message": "runs: {{ runs }}"})],
+)
+async def test_webhook_overide_body_job_failed(monkeypatch, override_body):
+    requests_mock = _mock_async_response(monkeypatch, "post", {"id": "response-id"})
+    runs = _generate_run_result(
+        state="error", error='can only concatenate str (not "int") to str'
+    )
+    await WebhookNotification(
+        params={"override_body": override_body, "url": "http://test.com"}
+    ).push("test-message", "info", [runs])
+    expected_body = {
+        "message": "runs: [{'project': 'test-remote-workflow', 'name': 'func-func', 'host': 'func-func-8lvl8', "
+        "'status': {'state': 'error', 'error': 'can only concatenate str (not \"int\") to str'}}]"
+    }
+    requests_mock.assert_called_once_with(
+        "http://test.com", headers={}, json=expected_body, ssl=None
+    )
+
+
 @pytest.mark.parametrize(
     "runs,expected,is_table",
     [
@@ -938,3 +979,63 @@ def _mock_async_response(monkeypatch, method, result):
     monkeypatch.setattr(aiohttp.ClientSession, method, requests_mock)
 
     return requests_mock
+
+
+def _generate_run_result(state: str, error: str = None, results: dict = None):
+    run_example = {
+        "status": {
+            "notifications": {
+                "Test": {"status": "pending", "sent_time": None, "reason": None}
+            },
+            "last_update": "2024-06-18T13:46:37.686443+00:00",
+            "start_time": "2024-06-18T13:46:37.392158+00:00",
+        },
+        "metadata": {
+            "uid": "b176e54e4ed24b28883aa69dce981601",
+            "project": "test-remote-workflow",
+            "name": "func-func",
+            "labels": {
+                "v3io_user": "admin",
+                "kind": "job",
+                "owner": "admin",
+                "mlrun/client_version": "1.7.0-rc21",
+                "mlrun/client_python_version": "3.9.18",
+                "host": "func-func-8lvl8",
+            },
+            "iteration": 0,
+        },
+        "spec": {
+            "function": "test-remote-workflow/func@8e0ddc3926470d5b97733679bb96738fa6dfd01b",
+            "parameters": {"x": 1},
+            "state_thresholds": {
+                "pending_scheduled": "1h",
+                "pending_not_scheduled": "-1",
+                "image_pull_backoff": "1h",
+                "executing": "24h",
+            },
+            "output_path": "v3io:///projects/test-remote-workflow/artifacts",
+            "notifications": [
+                {
+                    "when": ["error", "completed"],
+                    "name": "Test",
+                    "params": {
+                        "url": "https://webhook.site/5da7ac4d-39dc-4896-b18f-e13c5712a96a",
+                        "method": "POST",
+                    },
+                    "message": "",
+                    "status": "pending",
+                    "condition": "",
+                    "kind": "webhook",
+                    "severity": "info",
+                }
+            ],
+            "handler": "func",
+        },
+    }
+    if state == "completed":
+        run_example["status"]["results"] = results
+        run_example["status"]["state"] = state
+    elif state == "error":
+        run_example["status"]["error"] = error
+        run_example["status"]["state"] = state
+    return run_example