Allow custom wait delays (#4422)

# What this PR does Allows custom wait durations for: * `Wait` escalation policy * `>X alerts per Y minutes` escalation policy * `Wait` user notification policy ## Which issue(s) this PR closes Related to #2464 ## Checklist - [x] Unit, integration, and e2e (if applicable) tests updated - [x] Documentation added (or `pr:no public docs` PR label added if not required) - [x] Added the relevant release notes label (see labels prefixed w/ `release:`). These labels dictate how your PR will show up in the autogenerated release notes. --------- Co-authored-by: Rares Mardare <[email protected]>
grafana · May 31, 2024 · d8e1a1d · d8e1a1d
1 parent f7975f1
commit d8e1a1d
Show file tree

Hide file tree

Showing 19 changed files with 385 additions and 236 deletions.
diff --git a/engine/apps/api/serializers/escalation_policy.py b/engine/apps/api/serializers/escalation_policy.py
@@ -1,4 +1,3 @@
-import time
 from datetime import timedelta
 
 from rest_framework import serializers
@@ -9,6 +8,7 @@
 from apps.user_management.models import Team, User
 from apps.webhooks.models import Webhook
 from common.api_helpers.custom_fields import (
+    DurationSecondsField,
     OrganizationFilteredPrimaryKeyRelatedField,
     UsersFilteredByOrganizationField,
 )
@@ -47,15 +47,17 @@ class EscalationPolicySerializer(EagerLoadingMixin, serializers.ModelSerializer)
         queryset=User.objects,
         required=False,
     )
-    wait_delay = serializers.ChoiceField(
+    wait_delay = DurationSecondsField(
         required=False,
-        choices=EscalationPolicy.WEB_DURATION_CHOICES,
         allow_null=True,
+        min_value=timedelta(minutes=1),
+        max_value=timedelta(hours=24),
     )
-    num_minutes_in_window = serializers.ChoiceField(
+    num_minutes_in_window = serializers.IntegerField(
         required=False,
-        choices=EscalationPolicy.WEB_DURATION_CHOICES_MINUTES,
         allow_null=True,
+        min_value=1,  # 1 minute
+        max_value=24 * 60,  # 24 hours
     )
     notify_schedule = OrganizationFilteredPrimaryKeyRelatedField(
         queryset=OnCallSchedule.objects,
@@ -151,29 +153,12 @@ def validate_step(self, step_type):
             raise serializers.ValidationError("Invalid escalation step type: step is Slack-specific")
         return step_type
 
-    def to_internal_value(self, data):
-        data = self._wait_delay_to_internal_value(data)
-        return super().to_internal_value(data)
-
     def to_representation(self, instance):
         step = instance.step
         result = super().to_representation(instance)
         result = EscalationPolicySerializer._get_important_field(step, result)
         return result
 
-    @staticmethod
-    def _wait_delay_to_internal_value(data):
-        if data.get(WAIT_DELAY, None):
-            try:
-                time.strptime(data[WAIT_DELAY], "%H:%M:%S")
-            except ValueError:
-                try:
-                    data[WAIT_DELAY] = str(timedelta(seconds=float(data[WAIT_DELAY])))
-                except ValueError:
-                    raise serializers.ValidationError("Invalid wait delay format")
-
-        return data
-
     @staticmethod
     def _get_important_field(step, result):
         if step in {*EscalationPolicy.DEFAULT_STEPS_SET, *EscalationPolicy.STEPS_WITH_NO_IMPORTANT_VERSION_SET}:

diff --git a/engine/apps/api/serializers/user_notification_policy.py b/engine/apps/api/serializers/user_notification_policy.py
@@ -1,12 +1,11 @@
-import time
 from datetime import timedelta
 
 from rest_framework import serializers
 
 from apps.base.models import UserNotificationPolicy
 from apps.base.models.user_notification_policy import NotificationChannelAPIOptions
 from apps.user_management.models import User
-from common.api_helpers.custom_fields import OrganizationFilteredPrimaryKeyRelatedField
+from common.api_helpers.custom_fields import DurationSecondsField, OrganizationFilteredPrimaryKeyRelatedField
 from common.api_helpers.exceptions import Forbidden
 from common.api_helpers.mixins import EagerLoadingMixin
 
@@ -26,6 +25,12 @@ class UserNotificationPolicyBaseSerializer(EagerLoadingMixin, serializers.ModelS
         default=UserNotificationPolicy.Step.NOTIFY,
         choices=UserNotificationPolicy.Step.choices,
     )
+    wait_delay = DurationSecondsField(
+        required=False,
+        allow_null=True,
+        min_value=timedelta(minutes=1),
+        max_value=timedelta(hours=24),
+    )
 
     SELECT_RELATED = [
         "user",
@@ -41,14 +46,6 @@ class Meta:
         read_only_fields = ["order"]
 
     def to_internal_value(self, data):
-        if data.get("wait_delay", None):
-            try:
-                time.strptime(data["wait_delay"], "%H:%M:%S")
-            except ValueError:
-                try:
-                    data["wait_delay"] = str(timedelta(seconds=float(data["wait_delay"])))
-                except ValueError:
-                    raise serializers.ValidationError("Invalid wait delay format")
         data = self._notify_by_to_internal_value(data)
         return super().to_internal_value(data)
 

diff --git a/engine/apps/api/tests/test_escalation_policy.py b/engine/apps/api/tests/test_escalation_policy.py
@@ -39,7 +39,7 @@ def test_create_escalation_policy(escalation_policy_internal_api_setup, make_use
 
     data = {
         "step": EscalationPolicy.STEP_WAIT,
-        "wait_delay": "60.0",
+        "wait_delay": 60,
         "escalation_chain": escalation_chain.public_primary_key,
         "notify_to_users_queue": [],
         "from_time": None,
@@ -55,6 +55,28 @@ def test_create_escalation_policy(escalation_policy_internal_api_setup, make_use
     assert EscalationPolicy.objects.get(public_primary_key=response.data["id"]).order == max_order + 1
 
 
+@pytest.mark.django_db
+@pytest.mark.parametrize("wait_delay", (timedelta(seconds=59), timedelta(hours=24, seconds=1)))
+def test_create_escalation_policy_wait_delay_invalid(
+    escalation_policy_internal_api_setup, make_user_auth_headers, wait_delay
+):
+    token, escalation_chain, _, user, _ = escalation_policy_internal_api_setup
+    client = APIClient()
+    url = reverse("api-internal:escalation_policy-list")
+
+    data = {
+        "step": EscalationPolicy.STEP_WAIT,
+        "wait_delay": int(wait_delay.total_seconds()),
+        "escalation_chain": escalation_chain.public_primary_key,
+        "notify_to_users_queue": [],
+        "from_time": None,
+        "to_time": None,
+    }
+
+    response = client.post(url, data, format="json", **make_user_auth_headers(user, token))
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+
+
 @pytest.mark.django_db
 def test_create_escalation_policy_webhook(
     escalation_policy_internal_api_setup, make_custom_webhook, make_user_auth_headers
@@ -690,7 +712,7 @@ def test_escalation_policy_can_not_create_with_non_step_type_related_data(
         "escalation_chain": escalation_chain.public_primary_key,
         "step": step,
         "notify_to_users_queue": [user.public_primary_key],
-        "wait_delay": "300.0",
+        "wait_delay": 300,
         "from_time": "06:50:00",
         "to_time": "04:10:00",
     }

diff --git a/engine/apps/api/tests/test_user_notification_policy.py b/engine/apps/api/tests/test_user_notification_policy.py
@@ -1,4 +1,5 @@
 import json
+from datetime import timedelta
 from unittest.mock import patch
 
 import pytest
@@ -67,6 +68,26 @@ def test_create_notification_policy(user_notification_policy_internal_api_setup,
     assert response.status_code == status.HTTP_201_CREATED
 
 
+@pytest.mark.django_db
+@pytest.mark.parametrize("wait_delay", (timedelta(seconds=59), timedelta(hours=24, seconds=1)))
+def test_create_notification_policy_wait_delay_invalid(
+    user_notification_policy_internal_api_setup, make_user_auth_headers, wait_delay
+):
+    token, _, users = user_notification_policy_internal_api_setup
+    admin, _ = users
+    client = APIClient()
+    url = reverse("api-internal:notification_policy-list")
+
+    data = {
+        "step": UserNotificationPolicy.Step.WAIT,
+        "wait_delay": int(wait_delay.total_seconds()),
+        "important": False,
+        "user": admin.public_primary_key,
+    }
+    response = client.post(url, data, format="json", **make_user_auth_headers(admin, token))
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+
+
 @pytest.mark.django_db
 def test_admin_can_create_notification_policy_for_user(
     user_notification_policy_internal_api_setup, make_user_auth_headers
@@ -252,7 +273,7 @@ def test_unable_to_change_importance(user_notification_policy_internal_api_setup
 
 
 @pytest.mark.django_db
-@pytest.mark.parametrize("wait_delay, expected_wait_delay", [(None, "300.0"), ("900.0", "900.0")])
+@pytest.mark.parametrize("wait_delay, expected_wait_delay", [(None, 300), (900, 900)])
 def test_switch_step_type_from_notify_to_wait(
     make_organization_and_user_with_plugin_token,
     make_user_auth_headers,
@@ -400,9 +421,7 @@ def test_switch_notification_channel(
 
 
 @pytest.mark.django_db
-@pytest.mark.parametrize(
-    "from_wait_delay, to_wait_delay", [(None, "300.0"), (timezone.timedelta(seconds=900), "900.0")]
-)
+@pytest.mark.parametrize("from_wait_delay, to_wait_delay", [(None, 300), (timezone.timedelta(seconds=900), 900)])
 def test_switch_wait_delay(
     make_organization_and_user_with_plugin_token,
     make_user_auth_headers,

diff --git a/engine/apps/api/views/alert_group.py b/engine/apps/api/views/alert_group.py
@@ -643,6 +643,8 @@ def silence(self, request, pk=None):
     )
     @action(methods=["get"], detail=False)
     def silence_options(self, request):
+        # TODO: DEPRECATED, REMOVE IN A FUTURE RELEASE
+
         """
         Retrieve a list of valid silence options
         """

diff --git a/engine/apps/api/views/escalation_policy.py b/engine/apps/api/views/escalation_policy.py
@@ -138,13 +138,15 @@ def escalation_options(self, request):
 
     @action(detail=False, methods=["get"])
     def delay_options(self, request):
+        # TODO: DEPRECATED, REMOVE IN A FUTURE RELEASE
         choices = []
         for item in EscalationPolicy.WEB_DURATION_CHOICES:
             choices.append({"value": str(item[0]), "sec_value": item[0], "display_name": item[1]})
         return Response(choices)
 
     @action(detail=False, methods=["get"])
     def num_minutes_in_window_options(self, request):
+        # TODO: DEPRECATED, REMOVE IN A FUTURE RELEASE
         choices = [
             {"value": choice[0], "display_name": choice[1]} for choice in EscalationPolicy.WEB_DURATION_CHOICES_MINUTES
         ]

diff --git a/engine/common/api_helpers/custom_fields.py b/engine/common/api_helpers/custom_fields.py
@@ -1,3 +1,5 @@
+from datetime import timedelta
+
 from django.core.exceptions import ObjectDoesNotExist
 from drf_spectacular.utils import extend_schema_field
 from rest_framework import fields, serializers
@@ -205,3 +207,12 @@ def __init__(self, **kwargs):
             ],
             **kwargs,
         )
+
+
+# TODO: FloatField is used for backward-compatibility, change to IntegerField in a future release
+class DurationSecondsField(serializers.FloatField):
+    def to_internal_value(self, data):
+        return timedelta(seconds=int(super().to_internal_value(data)))
+
+    def to_representation(self, value):
+        return int(value.total_seconds())