From b41fec5439ca0d724c8b5b4d465feafee628e06a Mon Sep 17 00:00:00 2001 From: Matias Bordese Date: Mon, 5 Sep 2022 15:00:14 -0300 Subject: [PATCH] Add related escalation chains details to schedule API --- engine/apps/api/serializers/schedule_base.py | 7 ++ engine/apps/api/tests/test_schedules.py | 77 +++++++++++++++++++- engine/apps/api/views/schedule.py | 51 +++++++++---- 3 files changed, 117 insertions(+), 18 deletions(-) diff --git a/engine/apps/api/serializers/schedule_base.py b/engine/apps/api/serializers/schedule_base.py index 4cbd7641..b3e1858d 100644 --- a/engine/apps/api/serializers/schedule_base.py +++ b/engine/apps/api/serializers/schedule_base.py @@ -18,6 +18,7 @@ class ScheduleBaseSerializer(EagerLoadingMixin, serializers.ModelSerializer): user_group = UserGroupSerializer() warnings = serializers.SerializerMethodField() on_call_now = serializers.SerializerMethodField() + number_of_escalation_chains = serializers.SerializerMethodField() class Meta: fields = [ @@ -33,6 +34,7 @@ class ScheduleBaseSerializer(EagerLoadingMixin, serializers.ModelSerializer): "notify_empty_oncall", "mention_oncall_start", "mention_oncall_next", + "number_of_escalation_chains", ] SELECT_RELATED = ["organization"] @@ -71,6 +73,11 @@ class ScheduleBaseSerializer(EagerLoadingMixin, serializers.ModelSerializer): else: return [] + def get_number_of_escalation_chains(self, obj): + # num_escalation_chains param added in queryset via annotate. Check ScheduleView.get_queryset + # return 0 for just created schedules + return getattr(obj, "num_escalation_chains", 0) + def validate(self, attrs): if "slack_channel_id" in attrs: slack_channel_id = attrs.pop("slack_channel_id", None) diff --git a/engine/apps/api/tests/test_schedules.py b/engine/apps/api/tests/test_schedules.py index d2f90957..75fb2e52 100644 --- a/engine/apps/api/tests/test_schedules.py +++ b/engine/apps/api/tests/test_schedules.py @@ -9,6 +9,7 @@ from rest_framework.response import Response from rest_framework.serializers import ValidationError from rest_framework.test import APIClient +from apps.alerts.models import EscalationPolicy from apps.schedules.models import ( CustomOnCallShift, OnCallSchedule, @@ -57,11 +58,21 @@ def schedule_internal_api_setup( @pytest.mark.django_db -def test_get_list_schedules(schedule_internal_api_setup, make_user_auth_headers): +def test_get_list_schedules( + schedule_internal_api_setup, make_escalation_chain, make_escalation_policy, make_user_auth_headers +): user, token, calendar_schedule, ical_schedule, web_schedule, slack_channel = schedule_internal_api_setup client = APIClient() url = reverse("api-internal:schedule-list") + # setup escalation chain linked to web schedule + escalation_chain = make_escalation_chain(user.organization) + make_escalation_policy( + escalation_chain=escalation_chain, + escalation_policy_step=EscalationPolicy.STEP_NOTIFY_SCHEDULE, + notify_schedule=web_schedule, + ) + expected_payload = [ { "id": calendar_schedule.public_primary_key, @@ -79,6 +90,7 @@ def test_get_list_schedules(schedule_internal_api_setup, make_user_auth_headers) "mention_oncall_start": True, "notify_empty_oncall": 0, "notify_oncall_shift_freq": 1, + "number_of_escalation_chains": 0, }, { "id": ical_schedule.public_primary_key, @@ -96,6 +108,7 @@ def test_get_list_schedules(schedule_internal_api_setup, make_user_auth_headers) "mention_oncall_start": True, "notify_empty_oncall": 0, "notify_oncall_shift_freq": 1, + "number_of_escalation_chains": 0, }, { "id": web_schedule.public_primary_key, @@ -112,6 +125,7 @@ def test_get_list_schedules(schedule_internal_api_setup, make_user_auth_headers) "mention_oncall_start": True, "notify_empty_oncall": 0, "notify_oncall_shift_freq": 1, + "number_of_escalation_chains": 1, }, ] response = client.get(url, format="json", **make_user_auth_headers(user, token)) @@ -141,6 +155,7 @@ def test_get_detail_calendar_schedule(schedule_internal_api_setup, make_user_aut "mention_oncall_start": True, "notify_empty_oncall": 0, "notify_oncall_shift_freq": 1, + "number_of_escalation_chains": 0, } response = client.get(url, format="json", **make_user_auth_headers(user, token)) @@ -170,6 +185,7 @@ def test_get_detail_ical_schedule(schedule_internal_api_setup, make_user_auth_he "mention_oncall_start": True, "notify_empty_oncall": 0, "notify_oncall_shift_freq": 1, + "number_of_escalation_chains": 0, } response = client.get(url, format="json", **make_user_auth_headers(user, token)) @@ -178,8 +194,18 @@ def test_get_detail_ical_schedule(schedule_internal_api_setup, make_user_auth_he @pytest.mark.django_db -def test_get_detail_web_schedule(schedule_internal_api_setup, make_user_auth_headers): +def test_get_detail_web_schedule( + schedule_internal_api_setup, make_escalation_chain, make_escalation_policy, make_user_auth_headers +): user, token, _, _, web_schedule, _ = schedule_internal_api_setup + # setup escalation chain linked to web schedule + escalation_chain = make_escalation_chain(user.organization) + make_escalation_policy( + escalation_chain=escalation_chain, + escalation_policy_step=EscalationPolicy.STEP_NOTIFY_SCHEDULE, + notify_schedule=web_schedule, + ) + client = APIClient() url = reverse("api-internal:schedule-detail", kwargs={"pk": web_schedule.public_primary_key}) @@ -198,6 +224,7 @@ def test_get_detail_web_schedule(schedule_internal_api_setup, make_user_auth_hea "mention_oncall_start": True, "notify_empty_oncall": 0, "notify_oncall_shift_freq": 1, + "number_of_escalation_chains": 1, } response = client.get(url, format="json", **make_user_auth_headers(user, token)) @@ -230,6 +257,7 @@ def test_create_calendar_schedule(schedule_internal_api_setup, make_user_auth_he # modify initial data by adding id and None for optional fields schedule = OnCallSchedule.objects.get(public_primary_key=response.data["id"]) data["id"] = schedule.public_primary_key + data["number_of_escalation_chains"] = 0 assert response.status_code == status.HTTP_201_CREATED assert response.data == data @@ -262,6 +290,7 @@ def test_create_ical_schedule(schedule_internal_api_setup, make_user_auth_header # modify initial data by adding id and None for optional fields schedule = OnCallSchedule.objects.get(public_primary_key=response.data["id"]) data["id"] = schedule.public_primary_key + data["number_of_escalation_chains"] = 0 assert response.status_code == status.HTTP_201_CREATED assert response.data == data @@ -290,6 +319,7 @@ def test_create_web_schedule(schedule_internal_api_setup, make_user_auth_headers # modify initial data by adding id and None for optional fields schedule = OnCallSchedule.objects.get(public_primary_key=response.data["id"]) data["id"] = schedule.public_primary_key + data["number_of_escalation_chains"] = 0 assert response.status_code == status.HTTP_201_CREATED assert response.data == data @@ -868,6 +898,49 @@ def test_next_shifts_per_user( assert returned_data == expected +@pytest.mark.django_db +def test_related_escalation_chains( + make_organization_and_user_with_plugin_token, + make_user_auth_headers, + make_schedule, + make_escalation_chain, + make_escalation_policy, +): + organization, user, token = make_organization_and_user_with_plugin_token() + client = APIClient() + + schedule = make_schedule( + organization, + schedule_class=OnCallScheduleWeb, + name="test_web_schedule", + ) + # setup escalation chains linked to web schedule + escalation_chains = [] + for i in range(3): + chain = make_escalation_chain(user.organization) + make_escalation_policy( + escalation_chain=chain, + escalation_policy_step=EscalationPolicy.STEP_NOTIFY_SCHEDULE, + notify_schedule=schedule, + ) + escalation_chains.append(chain) + # setup other unrelated schedule + other_schedule = make_schedule(organization, schedule_class=OnCallScheduleWeb) + other_chain = make_escalation_chain(user.organization) + make_escalation_policy( + escalation_chain=other_chain, + escalation_policy_step=EscalationPolicy.STEP_NOTIFY_SCHEDULE, + notify_schedule=other_schedule, + ) + + url = reverse("api-internal:schedule-related-escalation-chains", kwargs={"pk": schedule.public_primary_key}) + response = client.get(url, format="json", **make_user_auth_headers(user, token)) + assert response.status_code == status.HTTP_200_OK + + expected = [{"name": chain.name, "pk": chain.public_primary_key} for chain in escalation_chains] + assert sorted(response.data, key=lambda e: e["name"]) == sorted(expected, key=lambda e: e["name"]) + + @pytest.mark.django_db def test_merging_same_shift_events( make_organization_and_user_with_plugin_token, diff --git a/engine/apps/api/views/schedule.py b/engine/apps/api/views/schedule.py index 7ef4792b..d992b302 100644 --- a/engine/apps/api/views/schedule.py +++ b/engine/apps/api/views/schedule.py @@ -1,6 +1,6 @@ import pytz from django.core.exceptions import ObjectDoesNotExist -from django.db.models import OuterRef, Subquery +from django.db.models import Count, OuterRef, Subquery from django.db.utils import IntegrityError from django.urls import reverse from django.utils import dateparse, timezone @@ -12,6 +12,7 @@ from rest_framework.permissions import IsAuthenticated from rest_framework.views import Response from rest_framework.viewsets import ModelViewSet +from apps.alerts.models import EscalationChain from apps.api.permissions import MODIFY_ACTIONS, READ_ACTIONS, ActionPermission, AnyRole, IsAdmin, IsAdminOrEditor from apps.api.serializers.schedule_base import ScheduleFastSerializer from apps.api.serializers.schedule_polymorphic import ( @@ -59,6 +60,7 @@ class ScheduleView( "notify_empty_oncall_options", "notify_oncall_shift_freq_options", "mention_options", + "related_escalation_chains", ), } @@ -90,6 +92,25 @@ class ScheduleView( context.update({"can_update_user_groups": self.can_update_user_groups}) return context + def _annotate_queryset(self, queryset): + """Annotate queryset with additional schedule metadata.""" + organization = self.request.auth.organization + slack_channels = SlackChannel.objects.filter( + slack_team_identity=organization.slack_team_identity, + slack_id=OuterRef("channel"), + ) + queryset = queryset.annotate( + slack_channel_name=Subquery(slack_channels.values("name")[:1]), + slack_channel_pk=Subquery(slack_channels.values("public_primary_key")[:1]), + ) + queryset = queryset.annotate( + num_escalation_chains=Count( + "escalation_policies__escalation_chain", + distinct=True, + ) + ) + return queryset + def get_queryset(self): is_short_request = self.request.query_params.get("short", "false") == "true" organization = self.request.auth.organization @@ -98,14 +119,7 @@ class ScheduleView( team=self.request.user.current_team, ) if not is_short_request: - slack_channels = SlackChannel.objects.filter( - slack_team_identity=organization.slack_team_identity, - slack_id=OuterRef("channel"), - ) - queryset = queryset.annotate( - slack_channel_name=Subquery(slack_channels.values("name")[:1]), - slack_channel_pk=Subquery(slack_channels.values("public_primary_key")[:1]), - ) + queryset = self._annotate_queryset(queryset) queryset = self.serializer_class.setup_eager_loading(queryset) return queryset @@ -113,14 +127,10 @@ class ScheduleView( # Override this method because we want to get object from organization instead of concrete team. pk = self.kwargs["pk"] organization = self.request.auth.organization - slack_channels = SlackChannel.objects.filter( - slack_team_identity=organization.slack_team_identity, - slack_id=OuterRef("channel"), - ) - queryset = organization.oncall_schedules.filter(public_primary_key=pk,).annotate( - slack_channel_name=Subquery(slack_channels.values("name")[:1]), - slack_channel_pk=Subquery(slack_channels.values("public_primary_key")[:1]), + queryset = organization.oncall_schedules.filter( + public_primary_key=pk, ) + queryset = self._annotate_queryset(queryset) try: obj = queryset.get() @@ -260,6 +270,15 @@ class ScheduleView( result = {"users": users} return Response(result, status=status.HTTP_200_OK) + @action(detail=True, methods=["get"]) + def related_escalation_chains(self, request, pk): + """Return escalation chains associated to schedule.""" + schedule = self.original_get_object() + escalation_chains = EscalationChain.objects.filter(escalation_policies__notify_schedule=schedule).distinct() + + result = [{"name": e.name, "pk": e.public_primary_key} for e in escalation_chains] + return Response(result, status=status.HTTP_200_OK) + @action(detail=False, methods=["get"]) def type_options(self, request): # TODO: check if it needed