diff --git a/engine/apps/api/tests/test_schedules.py b/engine/apps/api/tests/test_schedules.py index 4a29dc9d..8efcb6b2 100644 --- a/engine/apps/api/tests/test_schedules.py +++ b/engine/apps/api/tests/test_schedules.py @@ -1442,8 +1442,9 @@ def test_next_shifts_per_user( ("B", "UTC"), ("C", None), ("D", "America/Montevideo"), + ("E", None), ) - user_a, user_b, user_c, user_d = ( + user_a, user_b, user_c, user_d, user_e = ( make_user_for_organization(organization, username=i, _timezone=tz) for i, tz in users ) @@ -1469,8 +1470,7 @@ def test_next_shifts_per_user( ) on_call_shift.add_rolling_users([[user]]) - # override in the past: 17-18 / D - # won't be listed, but user D will still be included in the response + # override in the past, won't be listed: 17-18 / D override_data = { "start": tomorrow - timezone.timedelta(days=3), "rotation_start": tomorrow - timezone.timedelta(days=3), @@ -1483,6 +1483,7 @@ def test_next_shifts_per_user( override.add_rolling_users([[user_d]]) # override: 17-18 / C + # this is before C's shift, so it will be listed as upcoming override_data = { "start": tomorrow + timezone.timedelta(hours=17), "rotation_start": tomorrow + timezone.timedelta(hours=17), @@ -1494,11 +1495,26 @@ def test_next_shifts_per_user( ) override.add_rolling_users([[user_c]]) + # override: 17-18 / E + fifteend_days_later = tomorrow + timezone.timedelta(days=15) + override_data = { + "start": fifteend_days_later + timezone.timedelta(hours=17), + "rotation_start": fifteend_days_later + timezone.timedelta(hours=17), + "duration": timezone.timedelta(hours=1), + "schedule": schedule, + } + override = make_on_call_shift( + organization=organization, shift_type=CustomOnCallShift.TYPE_OVERRIDE, **override_data + ) + override.add_rolling_users([[user_e]]) + # final schedule: 7-12: B, 15-16: A, 16-17: B, 17-18: C (override), 18-20: C schedule.refresh_ical_final_schedule() url = reverse("api-internal:schedule-next-shifts-per-user", kwargs={"pk": schedule.public_primary_key}) - response = client.get(url, format="json", **make_user_auth_headers(admin, token)) + + # check for users with shifts in the next week + response = client.get(url + "?days=7", format="json", **make_user_auth_headers(admin, token)) assert response.status_code == status.HTTP_200_OK expected = { @@ -1517,13 +1533,27 @@ def test_next_shifts_per_user( tomorrow + timezone.timedelta(hours=18), user_c.timezone, ), - user_d.public_primary_key: (None, None, user_d.timezone), } returned_data = { u: (ev.get("start"), ev.get("end"), ev.get("user_timezone")) for u, ev in response.data["users"].items() } assert returned_data == expected + # by default it will check for shifts in the next 45 days + response = client.get(url, format="json", **make_user_auth_headers(admin, token)) + assert response.status_code == status.HTTP_200_OK + + # include user E with the override + expected[user_e.public_primary_key] = ( + fifteend_days_later + timezone.timedelta(hours=17), + fifteend_days_later + timezone.timedelta(hours=18), + user_e.timezone, + ) + returned_data = { + u: (ev.get("start"), ev.get("end"), ev.get("user_timezone")) for u, ev in response.data["users"].items() + } + assert returned_data == expected + @pytest.mark.django_db def test_next_shifts_per_user_ical_schedule_using_emails( diff --git a/engine/apps/api/views/schedule.py b/engine/apps/api/views/schedule.py index 78635290..e30aa8cb 100644 --- a/engine/apps/api/views/schedule.py +++ b/engine/apps/api/views/schedule.py @@ -388,20 +388,22 @@ class ScheduleView( @action(detail=True, methods=["get"]) def next_shifts_per_user(self, request, pk): """Return next shift for users in schedule.""" + days = self.request.query_params.get("days") + days = int(days) if days else 30 now = timezone.now() - datetime_end = now + datetime.timedelta(days=30) + datetime_end = now + datetime.timedelta(days=days) schedule = self.get_object(annotate=False) + users = {} events = schedule.final_events(now, datetime_end) - - # include user TZ information for every user - users = {u.public_primary_key: {"user_timezone": u.timezone} for u in schedule.related_users()} + users_tz = {u.public_primary_key: u.timezone for u in schedule.related_users()} added_users = set() for e in events: - user = e["users"][0]["pk"] if e["users"] else None - if user is not None and user not in added_users and user in users and e["end"] > now: - users[user].update(e) - added_users.add(user) + user_ppk = e["users"][0]["pk"] if e["users"] else None + if user_ppk is not None and user_ppk not in users and user_ppk in users_tz and e["end"] > now: + users[user_ppk] = e + users[user_ppk]["user_timezone"] = users_tz[user_ppk] + added_users.add(user_ppk) result = {"users": users} return Response(result, status=status.HTTP_200_OK)