Merge pull request #796 from grafana/matiasb/enforce-until-on-overrides

Enforce until on override shifts

engine/apps/api/views/schedule.py

@@ -1,5 +1,3 @@
-import datetime
-
 import pytz
 from django.core.exceptions import ObjectDoesNotExist
 from django.db.models import Count, OuterRef, Subquery
@@ -297,11 +295,7 @@ class ScheduleView(
         users = {u: None for u in schedule.related_users()}
         for e in events:
             user = e["users"][0]["pk"] if e["users"] else None
-            event_end = e["end"]
-            if not isinstance(event_end, datetime.datetime):
-                # all day events end is a date, make it a datetime for comparison
-                event_end = datetime.datetime.combine(event_end, datetime.datetime.min.time(), tzinfo=pytz.UTC)
-            if user is not None and users.get(user) is None and event_end > now:
+            if user is not None and users.get(user) is None and e["end"] > now:
                 users[user] = e
 
         result = {"users": users}

engine/apps/schedules/models/custom_on_call_shift.py

@@ -401,7 +401,10 @@ class CustomOnCallShift(models.Model):
         if user:
             event.add("summary", self.get_summary_with_user_for_ical(user))
         event.add("dtstart", self.convert_dt_to_schedule_timezone(start, time_zone))
-        event.add("dtend", self.convert_dt_to_schedule_timezone(start + self.duration, time_zone))
+        dtend = start + self.duration
+        if self.until:
+            dtend = min(dtend, self.until)
+        event.add("dtend", self.convert_dt_to_schedule_timezone(dtend, time_zone))
         event.add("dtstamp", self.rotation_start)
         if custom_rrule:
             event.add("rrule", custom_rrule)

engine/apps/schedules/models/on_call_schedule.py

@@ -272,11 +272,7 @@ class OnCallSchedule(PolymorphicModel):
             return []
 
         def event_start_cmp_key(e):
-            # all day events: compare using a datetime object at 00:00
-            start = e["start"]
-            if not isinstance(start, datetime.datetime):
-                start = datetime.datetime.combine(start, datetime.datetime.min.time(), tzinfo=pytz.UTC)
-            return start
+            return e["start"]
 
         def event_cmp_key(e):
             """Sorting key criteria for events."""

engine/apps/schedules/tests/test_custom_on_call_shift.py

@@ -59,6 +59,38 @@ def test_get_on_call_users_from_web_schedule_override(make_organization_and_user
     assert user in users_on_call
 
 
+@pytest.mark.django_db
+def test_get_on_call_users_from_web_schedule_override_until(
+    make_organization_and_user, make_on_call_shift, make_schedule
+):
+    organization, user = make_organization_and_user()
+
+    schedule = make_schedule(organization, schedule_class=OnCallScheduleWeb)
+    date = timezone.now().replace(microsecond=0)
+
+    data = {
+        "start": date,
+        "rotation_start": date,
+        "duration": timezone.timedelta(seconds=10800),
+        "schedule": schedule,
+        "until": date + timezone.timedelta(seconds=3600),
+    }
+
+    on_call_shift = make_on_call_shift(organization=organization, shift_type=CustomOnCallShift.TYPE_OVERRIDE, **data)
+    on_call_shift.add_rolling_users([[user]])
+
+    # user is on-call
+    date = date + timezone.timedelta(minutes=5)
+    users_on_call = list_users_to_notify_from_ical(schedule, date)
+    assert len(users_on_call) == 1
+    assert user in users_on_call
+
+    # and the until is enforced
+    date = date + timezone.timedelta(hours=2)
+    users_on_call = list_users_to_notify_from_ical(schedule, date)
+    assert len(users_on_call) == 0
+
+
 @pytest.mark.django_db
 def test_get_on_call_users_from_recurrent_event(make_organization_and_user, make_on_call_shift, make_schedule):
     organization, user = make_organization_and_user()