From 18fd3a4f71e292a3da4def7cc5565d20e3fadb2f Mon Sep 17 00:00:00 2001
From: Matias Bordese <mbordese@gmail.com>
Date: Fri, 19 Aug 2022 15:45:52 -0300
Subject: [PATCH] Update shift preview to reuse shift PK when previewing update

---
 engine/apps/api/tests/test_oncall_shift.py    | 20 ++++++++++++++++---
 engine/apps/api/views/on_call_shifts.py       |  3 +++
 .../apps/schedules/models/on_call_schedule.py |  6 ++++--
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/engine/apps/api/tests/test_oncall_shift.py b/engine/apps/api/tests/test_oncall_shift.py
index b31aa598..ff752555 100644
--- a/engine/apps/api/tests/test_oncall_shift.py
+++ b/engine/apps/api/tests/test_oncall_shift.py
@@ -1467,9 +1467,25 @@ def test_on_call_shift_preview_update(
 
     # check rotation events
     rotation_events = response.json()["rotation"]
+    # previewing an update reuses shift PK, so rotation keeps original event too
     expected_rotation_events = [
         {
             "calendar_type": OnCallSchedule.TYPE_ICAL_PRIMARY,
+            "shift": {"pk": on_call_shift.public_primary_key},
+            "start": on_call_shift.start.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "end": (on_call_shift.start + timezone.timedelta(hours=1)).strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "all_day": False,
+            "is_override": False,
+            "is_empty": False,
+            "is_gap": False,
+            "priority_level": 1,
+            "missing_users": [],
+            "users": [{"display_name": user.username, "pk": user.public_primary_key}],
+            "source": "api",
+        },
+        {
+            "calendar_type": OnCallSchedule.TYPE_ICAL_PRIMARY,
+            "shift": {"pk": on_call_shift.public_primary_key},
             "start": shift_start,
             "end": shift_end,
             "all_day": False,
@@ -1480,10 +1496,8 @@ def test_on_call_shift_preview_update(
             "missing_users": [],
             "users": [{"display_name": other_user.username, "pk": other_user.public_primary_key}],
             "source": "web",
-        }
+        },
     ]
-    # there isn't a saved shift, we don't care/know the temp pk
-    _ = [r.pop("shift") for r in rotation_events]
     assert rotation_events == expected_rotation_events
 
     # check final schedule events
diff --git a/engine/apps/api/views/on_call_shifts.py b/engine/apps/api/views/on_call_shifts.py
index 42c5eda2..9ca87d2f 100644
--- a/engine/apps/api/views/on_call_shifts.py
+++ b/engine/apps/api/views/on_call_shifts.py
@@ -90,6 +90,9 @@ class OnCallShiftView(PublicPrimaryKeyMixin, UpdateSerializerMixin, ModelViewSet
         validated_data = serializer._correct_validated_data(
             serializer.validated_data["type"], serializer.validated_data
         )
+        if not validated_data.get("rolling_users"):
+            return Response(data=serializer.errors, status=status.HTTP_400_BAD_REQUEST)
+
         updated_shift_pk = self.request.data.get("shift_pk")
         shift = CustomOnCallShift(**validated_data)
         schedule = shift.schedule
diff --git a/engine/apps/schedules/models/on_call_schedule.py b/engine/apps/schedules/models/on_call_schedule.py
index b89a8815..264b351d 100644
--- a/engine/apps/schedules/models/on_call_schedule.py
+++ b/engine/apps/schedules/models/on_call_schedule.py
@@ -631,9 +631,11 @@ class OnCallScheduleWeb(OnCallSchedule):
             except CustomOnCallShift.DoesNotExist:
                 pass
             else:
-                update_shift.until = custom_shift.rotation_start
+                if update_shift.event_is_started:
+                    update_shift.until = custom_shift.rotation_start
+                    extra_shifts.append(update_shift)
+                custom_shift.public_primary_key = updated_shift_pk
                 qs = qs.exclude(public_primary_key=updated_shift_pk)
-                extra_shifts.append(update_shift)
 
         ical_file = self._generate_ical_file_from_shifts(qs, extra_shifts=extra_shifts)