From 35620028ccbb88aa908f2f52627a226197662cda Mon Sep 17 00:00:00 2001
From: Matias Bordese <mbordese@gmail.com>
Date: Wed, 18 Oct 2023 14:14:56 -0300
Subject: [PATCH] Add user TZ information to next shifts per user endpoint
 (#3157)

---
 CHANGELOG.md                            |  4 ++++
 engine/apps/api/tests/test_schedules.py | 32 ++++++++++++++++++++-----
 engine/apps/api/views/schedule.py       |  9 ++++---
 3 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7e6e81ee..8a81dca4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Update ical schedule creation/update to trigger final schedule refresh ([#3156](https://github.com/grafana/oncall/pull/3156))
 
+### Changed
+
+- Add user TZ information to next shifts per user endpoint ([#3157](https://github.com/grafana/oncall/pull/3157))
+
 ## v1.3.44 (2023-10-16)
 
 ### Added
diff --git a/engine/apps/api/tests/test_schedules.py b/engine/apps/api/tests/test_schedules.py
index 756d8285..828aa80c 100644
--- a/engine/apps/api/tests/test_schedules.py
+++ b/engine/apps/api/tests/test_schedules.py
@@ -1398,7 +1398,15 @@ def test_next_shifts_per_user(
     )
 
     tomorrow = timezone.now().replace(hour=0, minute=0, second=0, microsecond=0) + timezone.timedelta(days=1)
-    user_a, user_b, user_c, user_d = (make_user_for_organization(organization, username=i) for i in "ABCD")
+    users = (
+        ("A", "Europe/London"),
+        ("B", "UTC"),
+        ("C", None),
+        ("D", "America/Montevideo"),
+    )
+    user_a, user_b, user_c, user_d = (
+        make_user_for_organization(organization, username=i, _timezone=tz) for i, tz in users
+    )
     # clear users pks <-> organization cache (persisting between tests)
     memoized_users_in_ical.cache_clear()
 
@@ -1456,13 +1464,25 @@ def test_next_shifts_per_user(
     assert response.status_code == status.HTTP_200_OK
 
     expected = {
-        user_a.public_primary_key: (tomorrow + timezone.timedelta(hours=15), tomorrow + timezone.timedelta(hours=16)),
-        user_b.public_primary_key: (tomorrow + timezone.timedelta(hours=7), tomorrow + timezone.timedelta(hours=12)),
-        user_c.public_primary_key: (tomorrow + timezone.timedelta(hours=17), tomorrow + timezone.timedelta(hours=18)),
-        user_d.public_primary_key: None,
+        user_a.public_primary_key: (
+            tomorrow + timezone.timedelta(hours=15),
+            tomorrow + timezone.timedelta(hours=16),
+            user_a.timezone,
+        ),
+        user_b.public_primary_key: (
+            tomorrow + timezone.timedelta(hours=7),
+            tomorrow + timezone.timedelta(hours=12),
+            user_b.timezone,
+        ),
+        user_c.public_primary_key: (
+            tomorrow + timezone.timedelta(hours=17),
+            tomorrow + timezone.timedelta(hours=18),
+            user_c.timezone,
+        ),
+        user_d.public_primary_key: (None, None, user_d.timezone),
     }
     returned_data = {
-        u: (ev["start"], ev["end"]) if ev is not None else None for u, ev in response.data["users"].items()
+        u: (ev.get("start"), ev.get("end"), ev.get("user_timezone")) for u, ev in response.data["users"].items()
     }
     assert returned_data == expected
 
diff --git a/engine/apps/api/views/schedule.py b/engine/apps/api/views/schedule.py
index eb300ee2..83121b1e 100644
--- a/engine/apps/api/views/schedule.py
+++ b/engine/apps/api/views/schedule.py
@@ -375,11 +375,14 @@ class ScheduleView(
 
         events = schedule.final_events(now, datetime_end)
 
-        users = {u.public_primary_key: None for u in schedule.related_users()}
+        # include user TZ information for every user
+        users = {u.public_primary_key: {"user_timezone": 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 users.get(user) is None and e["end"] > now:
-                users[user] = e
+            if user is not None and user not in added_users and e["end"] > now:
+                users[user].update(e)
+                added_users.add(user)
 
         result = {"users": users}
         return Response(result, status=status.HTTP_200_OK)