diff --git a/engine/apps/api/tests/test_shift_swaps.py b/engine/apps/api/tests/test_shift_swaps.py index e45ccb75..229fa1b1 100644 --- a/engine/apps/api/tests/test_shift_swaps.py +++ b/engine/apps/api/tests/test_shift_swaps.py @@ -176,8 +176,8 @@ def test_retrieve_permissions( assert response.status_code == expected_status -@patch("apps.api.views.shift_swap.write_resource_insight_log") -@patch("apps.api.views.shift_swap.create_shift_swap_request_message") +@patch("apps.schedules.models.shift_swap_request.write_resource_insight_log") +@patch("apps.schedules.tasks.shift_swaps.create_shift_swap_request_message") @pytest.mark.django_db def test_create( mock_create_shift_swap_request_message, diff --git a/engine/apps/api/views/shift_swap.py b/engine/apps/api/views/shift_swap.py index b6b94a71..6bad0b28 100644 --- a/engine/apps/api/views/shift_swap.py +++ b/engine/apps/api/views/shift_swap.py @@ -18,7 +18,7 @@ from apps.auth_token.auth import PluginAuthentication from apps.mobile_app.auth import MobileAppAuthTokenAuthentication from apps.schedules import exceptions from apps.schedules.models import ShiftSwapRequest -from apps.schedules.tasks.shift_swaps import create_shift_swap_request_message, update_shift_swap_request_message +from apps.schedules.tasks.shift_swaps import update_shift_swap_request_message from apps.user_management.models import User from common.api_helpers.exceptions import BadRequest from common.api_helpers.mixins import PublicPrimaryKeyMixin @@ -33,13 +33,6 @@ class BaseShiftSwapViewSet(ModelViewSet): serializer_class = ShiftSwapRequestSerializer pagination_class = FiftyPageSizePaginator - def _do_create(self, beneficiary: User, serializer: BaseSerializer[ShiftSwapRequest]) -> None: - shift_swap_request = serializer.save(beneficiary=beneficiary) - - write_resource_insight_log(instance=shift_swap_request, author=self.request.user, event=EntityEvent.CREATED) - - create_shift_swap_request_message.apply_async((shift_swap_request.pk,)) - def _do_take(self, benefactor: User) -> dict: shift_swap = self.get_object() prev_state = shift_swap.insight_logs_serialized @@ -83,7 +76,7 @@ class BaseShiftSwapViewSet(ModelViewSet): def perform_create(self, serializer: BaseSerializer[ShiftSwapRequest]) -> None: # default to create swap request with logged in user as beneficiary - self._do_create(self.request.user, serializer=serializer) + serializer.save(beneficiary=self.request.user) def perform_update(self, serializer: BaseSerializer[ShiftSwapRequest]) -> None: prev_state = serializer.instance.insight_logs_serialized diff --git a/engine/apps/google/client.py b/engine/apps/google/client.py index 3815fba7..906194ca 100644 --- a/engine/apps/google/client.py +++ b/engine/apps/google/client.py @@ -6,12 +6,23 @@ from django.conf import settings from google.oauth2.credentials import Credentials from googleapiclient.discovery import build -from apps.google.types import GoogleCalendarEvent +from apps.google import constants, utils +from apps.google.types import GoogleCalendarEvent as GoogleCalendarEventType logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) +class GoogleCalendarEvent: + def __init__(self, event: GoogleCalendarEventType): + self.raw_event = event + self._start_time = utils.datetime_strptime(event["start"]["dateTime"]) + self._end_time = utils.datetime_strptime(event["end"]["dateTime"]) + + self.start_time_utc = self._start_time.astimezone(datetime.timezone.utc) + self.end_time_utc = self._end_time.astimezone(datetime.timezone.utc) + + class GoogleCalendarAPIClient: MAX_NUMBER_OF_CALENDAR_EVENTS_TO_FETCH = 250 """ @@ -46,26 +57,23 @@ class GoogleCalendarAPIClient: """ https://developers.google.com/calendar/api/v3/reference/events/list """ - - def _format_datetime_arg(dt: datetime.datetime) -> str: - """ - https://stackoverflow.com/a/17159470/3902555 - """ - return dt.strftime("%Y-%m-%dT%H:%M:%SZ") - - now = _format_datetime_arg(datetime.datetime.now(datetime.UTC)) - logger.info( f"GoogleCalendarAPIClient - Getting the upcoming {self.MAX_NUMBER_OF_CALENDAR_EVENTS_TO_FETCH} " "out of office events" ) + now = datetime.datetime.now(datetime.UTC) + time_min = utils.datetime_strftime(now) + time_max = utils.datetime_strftime( + now + datetime.timedelta(days=constants.DAYS_IN_FUTURE_TO_CONSIDER_OUT_OF_OFFICE_EVENTS) + ) + events_result = ( self.service.events() .list( calendarId=self.CALENDAR_ID, - timeMin=now, - # timeMax= TODO: should we only fetch out of office events for next X amount of time? + timeMin=time_min, + timeMax=time_max, maxResults=self.MAX_NUMBER_OF_CALENDAR_EVENTS_TO_FETCH, singleEvents=True, orderBy="startTime", @@ -73,4 +81,4 @@ class GoogleCalendarAPIClient: ) .execute() ) - return events_result.get("items", []) + return [GoogleCalendarEvent(event) for event in events_result.get("items", [])] diff --git a/engine/apps/google/constants.py b/engine/apps/google/constants.py new file mode 100644 index 00000000..8d91becf --- /dev/null +++ b/engine/apps/google/constants.py @@ -0,0 +1,6 @@ +GOOGLE_CALENDAR_EVENT_DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S%z" +""" +https://stackoverflow.com/a/17159470/3902555 +""" + +DAYS_IN_FUTURE_TO_CONSIDER_OUT_OF_OFFICE_EVENTS = 30 diff --git a/engine/apps/google/models/google_oauth2_user.py b/engine/apps/google/models/google_oauth2_user.py index 268acbe4..aa03a731 100644 --- a/engine/apps/google/models/google_oauth2_user.py +++ b/engine/apps/google/models/google_oauth2_user.py @@ -1,8 +1,15 @@ +import typing + from django.db import models from mirage import fields as mirage_fields +if typing.TYPE_CHECKING: + from apps.user_management.models import User + class GoogleOAuth2User(models.Model): + user: "User" + id = models.AutoField(primary_key=True) user = models.OneToOneField( to="user_management.User", null=False, blank=False, related_name="google_oauth2_user", on_delete=models.CASCADE diff --git a/engine/apps/google/tasks.py b/engine/apps/google/tasks.py index 4572dd91..f8af0375 100644 --- a/engine/apps/google/tasks.py +++ b/engine/apps/google/tasks.py @@ -4,6 +4,7 @@ from celery.utils.log import get_task_logger from apps.google.client import GoogleCalendarAPIClient from apps.google.models import GoogleOAuth2User +from apps.schedules.models import OnCallSchedule, ShiftSwapRequest from common.custom_celery_tasks import shared_dedicated_queue_retry_task logger = get_task_logger(__name__) @@ -15,9 +16,69 @@ def sync_out_of_office_calendar_events_for_user(google_oauth2_user_pk: int) -> N google_oauth2_user = GoogleOAuth2User.objects.get(pk=google_oauth2_user_pk) google_api_client = GoogleCalendarAPIClient(google_oauth2_user.access_token, google_oauth2_user.refresh_token) - # NOTE: shift swap request generation will be done in https://github.com/grafana/oncall-private/issues/2590 - # QUESTION: will we need to persist any information about these calendar events in our database? - google_api_client.fetch_out_of_office_events() + user = google_oauth2_user.user + user_id = user.pk + + logger.info(f"Syncing out of office Google Calendar events for user {user_id}") + + users_schedules = OnCallSchedule.objects.related_to_user(user) + user_google_calendar_settings = user.google_calendar_settings + oncall_schedules_to_consider_for_shift_swaps = user_google_calendar_settings[ + "oncall_schedules_to_consider_for_shift_swaps" + ] + + if oncall_schedules_to_consider_for_shift_swaps: + users_schedules = users_schedules.filter(public_primary_key__in=oncall_schedules_to_consider_for_shift_swaps) + + for out_of_office_event in google_api_client.fetch_out_of_office_events(): + event_id = out_of_office_event.raw_event["id"] + start_time_utc = out_of_office_event.start_time_utc + end_time_utc = out_of_office_event.end_time_utc + + logger.info( + f"Processing out of office event {event_id} starting at {start_time_utc} and ending at " + f"{end_time_utc} for user {user_id}" + ) + + for schedule in users_schedules: + _, _, upcoming_shifts = schedule.shifts_for_user( + user, + start_time_utc, + datetime_end=end_time_utc, + ) + + if upcoming_shifts: + logger.info( + f"Found {len(upcoming_shifts)} upcoming shift(s) for user {user_id} " + f"during the out of office event {event_id}" + ) + + shift_swap_request_exists = ShiftSwapRequest.objects.filter( + beneficiary=user, + schedule=schedule, + swap_start=start_time_utc, + swap_end=end_time_utc, + ).exists() + + if not shift_swap_request_exists: + logger.info( + f"Creating shift swap request for user {user_id} schedule {schedule.pk} " + f"due to the out of office event {event_id}" + ) + + ssr = ShiftSwapRequest.objects.create( + beneficiary=user, + schedule=schedule, + swap_start=start_time_utc, + swap_end=end_time_utc, + description=f"{user.name or user.email} will be out of office during this time according to Google Calendar", + ) + + logger.info(f"Created shift swap request {ssr.pk}") + else: + logger.info(f"Shift swap request already exists for user {user_id} schedule {schedule.pk}") + else: + logger.info(f"No upcoming shifts found for user {user_id} during the out of office event {event_id}") @shared_dedicated_queue_retry_task(autoretry_for=(Exception,), retry_backoff=True) diff --git a/engine/apps/google/tests/__init__.py b/engine/apps/google/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/engine/apps/google/tests/factories.py b/engine/apps/google/tests/factories.py new file mode 100644 index 00000000..db54200e --- /dev/null +++ b/engine/apps/google/tests/factories.py @@ -0,0 +1,14 @@ +import factory + +from apps.google.models import GoogleOAuth2User +from common.utils import UniqueFaker + + +class GoogleOAuth2UserFactory(factory.DjangoModelFactory): + google_user_id = UniqueFaker("pyint") + access_token = factory.Faker("password") + refresh_token = factory.Faker("password") + oauth_scope = factory.Faker("word") + + class Meta: + model = GoogleOAuth2User diff --git a/engine/apps/google/tests/test_sync_out_of_office_calendar_events_for_user.py b/engine/apps/google/tests/test_sync_out_of_office_calendar_events_for_user.py new file mode 100644 index 00000000..0ad6f50b --- /dev/null +++ b/engine/apps/google/tests/test_sync_out_of_office_calendar_events_for_user.py @@ -0,0 +1,291 @@ +import datetime +from unittest.mock import patch + +import pytest +from django.utils import timezone + +from apps.google import constants, tasks +from apps.schedules.models import CustomOnCallShift, OnCallScheduleWeb, ShiftSwapRequest + + +def _create_mock_google_calendar_event(start_time: datetime.datetime, end_time: datetime.datetime): + return { + "colorId": "4", + "created": "2024-03-22T23:06:39.000Z", + "creator": { + "email": "joey.orlando@grafana.com", + "self": True, + }, + "end": { + "dateTime": end_time.strftime(constants.GOOGLE_CALENDAR_EVENT_DATETIME_FORMAT), + "timeZone": "America/New_York", + }, + "etag": "3422297608598000", + "eventType": "outOfOffice", + "extendedProperties": { + "private": { + "reclaim.event.category": "VACATION", + "reclaim.priority.index": "3", + "reclaim.project.id": "NULL", + "reclaim.touched": "true", + }, + }, + "htmlLink": "https://www.google.com/calendar/event?eid=NDlyZGVmNHU2aTVkaDR1aWFycGZqYWoya3Qgam9leS5vcmxhbmRvQGdyYWZhbmEuY29t", + "iCalUID": "49rdef4u6i5dh4uiarpfjaj2kt@google.com", + "id": "49rdef4u6i5dh4uiarpfjaj2kt", + "kind": "calendar#event", + "organizer": { + "email": "joey.orlando@grafana.com", + "self": True, + }, + "outOfOfficeProperties": { + "autoDeclineMode": "declineNone", + }, + "reminders": { + "useDefault": False, + }, + "sequence": 0, + "start": { + "dateTime": start_time.strftime(constants.GOOGLE_CALENDAR_EVENT_DATETIME_FORMAT), + "timeZone": "America/New_York", + }, + "status": "confirmed", + "summary": "Out of office", + "updated": "2024-03-22T23:06:44.299Z", + "visibility": "public", + } + + +def _create_event_start_and_end_times(start_days_in_future=5, end_time_minutes_past_start=50): + start_time = ( + datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(days=start_days_in_future) + ).replace(second=0, microsecond=0) + end_time = start_time + datetime.timedelta(minutes=end_time_minutes_past_start) + + return start_time, end_time + + +@pytest.fixture +def make_schedule_with_on_call_shift(make_schedule, make_on_call_shift): + def _make_schedule_with_on_call_shift(out_of_office_events, organization, user): + schedule = make_schedule( + organization, + schedule_class=OnCallScheduleWeb, + channel="channel", + prev_ical_file_overrides=None, + cached_ical_file_overrides=None, + ) + + dt_format = constants.GOOGLE_CALENDAR_EVENT_DATETIME_FORMAT + + if out_of_office_events: + on_call_shift_start = datetime.datetime.strptime( + out_of_office_events[0]["start"]["dateTime"], dt_format + ) - datetime.timedelta(days=60) + else: + on_call_shift_start = timezone.now() - datetime.timedelta(days=60) + + on_call_shift = make_on_call_shift( + organization=organization, + shift_type=CustomOnCallShift.TYPE_ROLLING_USERS_EVENT, + start=on_call_shift_start, + rotation_start=on_call_shift_start, + duration=datetime.timedelta(days=365), + priority_level=1, + frequency=CustomOnCallShift.FREQUENCY_DAILY, + schedule=schedule, + ) + on_call_shift.add_rolling_users([[user]]) + schedule.refresh_ical_file() + schedule.refresh_ical_final_schedule() + + return schedule + + return _make_schedule_with_on_call_shift + + +@pytest.fixture +def test_setup( + make_organization, + make_user_for_organization, + make_google_oauth2_user_for_user, + make_schedule_with_on_call_shift, +): + def _test_setup(out_of_office_events): + organization = make_organization() + user_name = "Bob Smith" + user = make_user_for_organization( + organization, + # normally this 👇 is done via User.finish_google_oauth2_connection_flow.. but since we're creating + # the user via a fixture we need to manually add this + google_calendar_settings={ + "oncall_schedules_to_consider_for_shift_swaps": [], + }, + name=user_name, + ) + + google_oauth2_user = make_google_oauth2_user_for_user(user) + schedule = make_schedule_with_on_call_shift(out_of_office_events, organization, user) + + return google_oauth2_user, schedule + + return _test_setup + + +@patch("apps.google.client.build") +@pytest.mark.django_db +def test_sync_out_of_office_calendar_events_for_user_no_ooo_events(mock_google_api_client_build, test_setup): + out_of_office_events = [] + + mock_google_api_client_build.return_value.events.return_value.list.return_value.execute.return_value = { + "items": out_of_office_events, + } + + google_oauth2_user, schedule = test_setup(out_of_office_events) + user = google_oauth2_user.user + + tasks.sync_out_of_office_calendar_events_for_user(google_oauth2_user.pk) + + assert ShiftSwapRequest.objects.filter(beneficiary=user, schedule=schedule).count() == 0 + + +@patch("apps.google.client.build") +@pytest.mark.django_db +def test_sync_out_of_office_calendar_events_for_user_single_ooo_event(mock_google_api_client_build, test_setup): + start_time, end_time = _create_event_start_and_end_times() + out_of_office_events = [ + _create_mock_google_calendar_event(start_time, end_time), + ] + + mock_google_api_client_build.return_value.events.return_value.list.return_value.execute.return_value = { + "items": out_of_office_events, + } + + google_oauth2_user, schedule = test_setup(out_of_office_events) + user = google_oauth2_user.user + + tasks.sync_out_of_office_calendar_events_for_user(google_oauth2_user.pk) + + ssrs = ShiftSwapRequest.objects.filter(beneficiary=user, schedule=schedule) + ssr = ssrs.first() + + assert ssrs.count() == 1 + + assert ssr.swap_start == start_time + assert ssr.swap_end == end_time + assert ssr.description == f"{user.name} will be out of office during this time according to Google Calendar" + + +@patch("apps.google.client.build") +@pytest.mark.django_db +def test_sync_out_of_office_calendar_events_for_user_multiple_ooo_events(mock_google_api_client_build, test_setup): + # partial day out of office event + event1_start_time, event1_end_time = _create_event_start_and_end_times() + # all day out of office event + event2_start_time, event2_end_time = _create_event_start_and_end_times(6, 24 * 60) + + out_of_office_events = [ + _create_mock_google_calendar_event(event1_start_time, event1_end_time), + _create_mock_google_calendar_event(event2_start_time, event2_end_time), + ] + + mock_google_api_client_build.return_value.events.return_value.list.return_value.execute.return_value = { + "items": out_of_office_events, + } + + google_oauth2_user, schedule = test_setup(out_of_office_events) + user = google_oauth2_user.user + + tasks.sync_out_of_office_calendar_events_for_user(google_oauth2_user.pk) + + assert ShiftSwapRequest.objects.filter(beneficiary=user, schedule=schedule).count() == 2 + + +@patch("apps.google.client.build") +@pytest.mark.django_db +def test_sync_out_of_office_calendar_events_for_user_oncall_schedules_to_consider_for_shift_swaps_setting( + mock_google_api_client_build, + test_setup, + make_schedule_with_on_call_shift, +): + start_time, end_time = _create_event_start_and_end_times() + out_of_office_events = [ + _create_mock_google_calendar_event(start_time, end_time), + ] + + mock_google_api_client_build.return_value.events.return_value.list.return_value.execute.return_value = { + "items": out_of_office_events, + } + + google_oauth2_user, schedule1 = test_setup(out_of_office_events) + user = google_oauth2_user.user + make_schedule_with_on_call_shift(out_of_office_events, schedule1.organization, user) + + user.google_calendar_settings = { + "oncall_schedules_to_consider_for_shift_swaps": [schedule1.public_primary_key], + } + user.save() + + tasks.sync_out_of_office_calendar_events_for_user(google_oauth2_user.pk) + + assert ShiftSwapRequest.objects.filter(beneficiary=user).count() == 1 + ssr = ShiftSwapRequest.objects.first() + + assert ssr.schedule == schedule1 + + +@patch("apps.google.tasks.OnCallSchedule.shifts_for_user", return_value=([], [], [])) +@patch("apps.google.client.build") +@pytest.mark.django_db +def test_sync_out_of_office_calendar_events_for_user_no_upcoming_shifts( + mock_google_api_client_build, + _mock_schedule_shifts_for_user, + test_setup, +): + start_time, end_time = _create_event_start_and_end_times() + out_of_office_events = [ + _create_mock_google_calendar_event(start_time, end_time), + ] + + mock_google_api_client_build.return_value.events.return_value.list.return_value.execute.return_value = { + "items": out_of_office_events, + } + + google_oauth2_user, _ = test_setup(out_of_office_events) + user = google_oauth2_user.user + + tasks.sync_out_of_office_calendar_events_for_user(google_oauth2_user.pk) + + assert ShiftSwapRequest.objects.filter(beneficiary=user).count() == 0 + + +@patch("apps.google.client.build") +@pytest.mark.django_db +def test_sync_out_of_office_calendar_events_for_user_preexisting_shift_swap_request( + mock_google_api_client_build, + test_setup, + make_shift_swap_request, +): + start_time, end_time = _create_event_start_and_end_times() + out_of_office_events = [ + _create_mock_google_calendar_event(start_time, end_time), + ] + + mock_google_api_client_build.return_value.events.return_value.list.return_value.execute.return_value = { + "items": out_of_office_events, + } + + google_oauth2_user, schedule = test_setup(out_of_office_events) + user = google_oauth2_user.user + + make_shift_swap_request( + schedule, + user, + swap_start=start_time, + swap_end=end_time, + ) + + tasks.sync_out_of_office_calendar_events_for_user(google_oauth2_user.pk) + + # should be 1 because we just created a shift swap request above via the fixture + assert ShiftSwapRequest.objects.filter(beneficiary=user, schedule=schedule).count() == 1 diff --git a/engine/apps/google/types.py b/engine/apps/google/types.py index fb431fe4..f62259e8 100644 --- a/engine/apps/google/types.py +++ b/engine/apps/google/types.py @@ -2,10 +2,15 @@ import typing class GoogleCalendarEventDate(typing.TypedDict): - date: typing.NotRequired[str] - """ - The date, in the format "yyyy-mm-dd", if this is an all-day event. - """ + # NOTE: in reality I haven't seen this field returned, even despite creating + # an out of office event with the "All day" checkbox checked. Instead it looks + # like it just returns the start.dateTime and end.dateTime as midnight of the + # respective days + + # date: typing.NotRequired[str] + # """ + # The date, in the format "yyyy-mm-dd", if this is an all-day event. + # """ dateTime: typing.NotRequired[str] """ diff --git a/engine/apps/google/utils.py b/engine/apps/google/utils.py new file mode 100644 index 00000000..ce30cd99 --- /dev/null +++ b/engine/apps/google/utils.py @@ -0,0 +1,11 @@ +import datetime + +from apps.google import constants + + +def datetime_strftime(dt: datetime.datetime) -> str: + return dt.strftime(constants.GOOGLE_CALENDAR_EVENT_DATETIME_FORMAT) + + +def datetime_strptime(dt: str) -> datetime.datetime: + return datetime.datetime.strptime(dt, constants.GOOGLE_CALENDAR_EVENT_DATETIME_FORMAT) diff --git a/engine/apps/public_api/tests/test_shift_swap.py b/engine/apps/public_api/tests/test_shift_swap.py index 44ca053b..c4828d11 100644 --- a/engine/apps/public_api/tests/test_shift_swap.py +++ b/engine/apps/public_api/tests/test_shift_swap.py @@ -112,8 +112,8 @@ def test_list_filters( assert_expected(response, (swap4,)) -@patch("apps.api.views.shift_swap.write_resource_insight_log") -@patch("apps.api.views.shift_swap.create_shift_swap_request_message") +@patch("apps.schedules.models.shift_swap_request.write_resource_insight_log") +@patch("apps.schedules.tasks.shift_swaps.create_shift_swap_request_message") @pytest.mark.django_db def test_create( mock_create_shift_swap_request_message, @@ -122,7 +122,7 @@ def test_create( make_user_for_organization, make_schedule, ): - organization, user, token = make_organization_and_user_with_token() + organization, _, token = make_organization_and_user_with_token() another_user = make_user_for_organization(organization) schedule = make_schedule(organization, schedule_class=OnCallScheduleWeb) today = timezone.now().replace(hour=0, minute=0, second=0, microsecond=0) @@ -145,7 +145,9 @@ def test_create( assert_swap_response(response, data) ssr = ShiftSwapRequest.objects.get(public_primary_key=response.json()["id"]) - mock_write_resource_insight_log.assert_called_once_with(instance=ssr, author=user, event=EntityEvent.CREATED) + mock_write_resource_insight_log.assert_called_once_with( + instance=ssr, author=another_user, event=EntityEvent.CREATED + ) mock_create_shift_swap_request_message.apply_async.assert_called_once_with((ssr.pk,)) diff --git a/engine/apps/public_api/views/shift_swap.py b/engine/apps/public_api/views/shift_swap.py index 73c725b1..29d5fcbe 100644 --- a/engine/apps/public_api/views/shift_swap.py +++ b/engine/apps/public_api/views/shift_swap.py @@ -85,8 +85,7 @@ class ShiftSwapViewSet(RateLimitHeadersMixin, BaseShiftSwapViewSet): return user def perform_create(self, serializer: BaseSerializer[ShiftSwapRequest]) -> None: - beneficiary = self._get_user("beneficiary") - self._do_create(beneficiary=beneficiary, serializer=serializer) + serializer.save(beneficiary=self._get_user("beneficiary")) @action(methods=["post"], detail=True) def take(self, request: AuthenticatedRequest, pk: str) -> Response: diff --git a/engine/apps/schedules/models/on_call_schedule.py b/engine/apps/schedules/models/on_call_schedule.py index 34c4f325..afa99664 100644 --- a/engine/apps/schedules/models/on_call_schedule.py +++ b/engine/apps/schedules/models/on_call_schedule.py @@ -542,10 +542,23 @@ class OnCallSchedule(PolymorphicModel): self.save(update_fields=["cached_ical_final_schedule"]) def shifts_for_user( - self, user: User, datetime_start: datetime.datetime, days: int = 7 + self, + user: User, + datetime_start: datetime.datetime, + datetime_end: typing.Optional[datetime.datetime] = None, + days: typing.Optional[int] = None, ) -> typing.Tuple[ScheduleEvents, ScheduleEvents, ScheduleEvents]: + """ + NOTE: must specify at least `datetime_end` or `days` + """ + if not datetime_end and not days: + raise ValueError("Must specify at least `datetime_end` or `days`") + now = timezone.now() - datetime_end = datetime_start + datetime.timedelta(days=days) + + if days is not None: + datetime_end = datetime_start + datetime.timedelta(days=days) + passed_shifts: ScheduleEvents = [] current_shifts: ScheduleEvents = [] upcoming_shifts: ScheduleEvents = [] diff --git a/engine/apps/schedules/models/shift_swap_request.py b/engine/apps/schedules/models/shift_swap_request.py index 586a4950..dde99a7f 100644 --- a/engine/apps/schedules/models/shift_swap_request.py +++ b/engine/apps/schedules/models/shift_swap_request.py @@ -5,10 +5,13 @@ from django.conf import settings from django.core.validators import MinLengthValidator from django.db import models from django.db.models import QuerySet +from django.db.models.signals import post_save +from django.dispatch import receiver from django.utils import timezone from apps.schedules import exceptions from apps.schedules.tasks import refresh_ical_final_schedule +from common.insight_log import EntityEvent, write_resource_insight_log from common.public_primary_keys import generate_public_primary_key, increase_public_primary_key_length if typing.TYPE_CHECKING: @@ -256,3 +259,14 @@ class ShiftSwapRequest(models.Model): result["schedule"] = self.schedule.insight_logs_verbal result["schedule_id"] = self.schedule.public_primary_key return result + + +@receiver(post_save, sender=ShiftSwapRequest) +def listen_for_shiftswaprequest_model_save( + sender: ShiftSwapRequest, instance: ShiftSwapRequest, created: bool, *args, **kwargs +) -> None: + from apps.schedules.tasks.shift_swaps import create_shift_swap_request_message + + if created: + write_resource_insight_log(instance=instance, author=instance.beneficiary, event=EntityEvent.CREATED) + create_shift_swap_request_message.apply_async((instance.pk,)) diff --git a/engine/apps/schedules/tests/test_on_call_schedule.py b/engine/apps/schedules/tests/test_on_call_schedule.py index 5fc66e49..6baa33fd 100644 --- a/engine/apps/schedules/tests/test_on_call_schedule.py +++ b/engine/apps/schedules/tests/test_on_call_schedule.py @@ -2663,7 +2663,7 @@ def test_shifts_for_user( schedule.refresh_ical_file() schedule.refresh_ical_final_schedule() - passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(admin, now) + passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(admin, now, days=7) assert len(passed_shifts) == 0 assert len(current_shifts) == 1 assert len(upcoming_shifts) == 7 @@ -2678,7 +2678,7 @@ def test_shifts_for_user( users = {u["pk"] for u in shift["users"]} assert admin.public_primary_key in users - passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(other_user, now) + passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(other_user, now, days=7) assert len(passed_shifts) == 0 assert len(current_shifts) == 0 assert len(upcoming_shifts) == 0 @@ -2731,7 +2731,7 @@ def test_shifts_for_user_only_two_users_with_shifts( schedule.refresh_ical_final_schedule() - passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(current_user, start_date, days) + passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(current_user, start_date, days=days) assert len(passed_shifts) == 0 assert len(current_shifts) == 0 assert len(upcoming_shifts) == 4 @@ -2740,7 +2740,7 @@ def test_shifts_for_user_only_two_users_with_shifts( assert current_user.public_primary_key in users assert shift["start"] > now - passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(user2, start_date, days) + passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(user2, start_date, days=days) assert len(passed_shifts) > 0 assert len(current_shifts) > 0 assert len(upcoming_shifts) > 0 @@ -2774,7 +2774,7 @@ def test_shifts_for_user_no_events( start_date = today - timezone.timedelta(days=2) days = 7 - passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(current_user, start_date, days) + passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(current_user, start_date, days=days) assert len(passed_shifts) == 0 assert len(current_shifts) == 0 assert len(upcoming_shifts) == 0 @@ -2795,7 +2795,7 @@ def test_shifts_for_user_without_final_ical( start_date = today - timezone.timedelta(days=2) days = 7 - passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(user, start_date, days) + passed_shifts, current_shifts, upcoming_shifts = schedule.shifts_for_user(user, start_date, days=days) assert len(passed_shifts) == 0 assert len(current_shifts) == 0 assert len(upcoming_shifts) == 0 diff --git a/engine/apps/user_management/tests/test_user.py b/engine/apps/user_management/tests/test_user.py index 09fa280a..7dc8a384 100644 --- a/engine/apps/user_management/tests/test_user.py +++ b/engine/apps/user_management/tests/test_user.py @@ -4,6 +4,7 @@ import pytest from django.utils import timezone from apps.api.permissions import LegacyAccessControlRole +from apps.google.models import GoogleOAuth2User from apps.user_management.models import User @@ -108,3 +109,77 @@ def test_is_telegram_connected(make_organization_and_user, make_telegram_user_co assert user.is_telegram_connected is False make_telegram_user_connector(user) assert user.is_telegram_connected is True + + +@pytest.mark.django_db +def test_has_google_oauth2_connected(make_organization_and_user, make_google_oauth2_user_for_user): + _, user = make_organization_and_user() + + assert user.has_google_oauth2_connected is False + make_google_oauth2_user_for_user(user) + assert user.has_google_oauth2_connected is True + + +@pytest.mark.django_db +def test_finish_google_oauth2_connection_flow(make_organization_and_user): + oauth_response = { + "access_token": "access", + "refresh_token": "refresh", + "sub": "google_user_id", + "scope": "scope", + } + + _, user = make_organization_and_user() + + assert GoogleOAuth2User.objects.filter(user=user).exists() is False + assert user.google_calendar_settings is None + + user.finish_google_oauth2_connection_flow(oauth_response) + user.refresh_from_db() + + google_oauth_user = user.google_oauth2_user + assert google_oauth_user.google_user_id == "google_user_id" + assert google_oauth_user.access_token == "access" + assert google_oauth_user.refresh_token == "refresh" + assert google_oauth_user.oauth_scope == "scope" + assert user.google_calendar_settings["oncall_schedules_to_consider_for_shift_swaps"] == [] + + oauth_response2 = { + "access_token": "access2", + "refresh_token": "refresh2", + "sub": "google_user_id2", + "scope": "scope2", + } + + user.finish_google_oauth2_connection_flow(oauth_response2) + user.refresh_from_db() + + google_oauth_user = user.google_oauth2_user + assert google_oauth_user.google_user_id == "google_user_id2" + assert google_oauth_user.access_token == "access2" + assert google_oauth_user.refresh_token == "refresh2" + assert google_oauth_user.oauth_scope == "scope2" + + +@pytest.mark.django_db +def test_finish_google_oauth2_disconnection_flow(make_organization_and_user): + _, user = make_organization_and_user() + + user.finish_google_oauth2_connection_flow( + { + "access_token": "access", + "refresh_token": "refresh", + "sub": "google_user_id", + "scope": "scope", + } + ) + user.refresh_from_db() + + assert user.google_oauth2_user is not None + assert user.google_calendar_settings is not None + + user.finish_google_oauth2_disconnection_flow() + user.refresh_from_db() + + assert GoogleOAuth2User.objects.filter(user=user).exists() is False + assert user.google_calendar_settings is None diff --git a/engine/conftest.py b/engine/conftest.py index ad222572..e461d670 100644 --- a/engine/conftest.py +++ b/engine/conftest.py @@ -62,6 +62,7 @@ from apps.base.tests.factories import ( UserNotificationPolicyLogRecordFactory, ) from apps.email.tests.factories import EmailMessageFactory +from apps.google.tests.factories import GoogleOAuth2UserFactory from apps.heartbeat.tests.factories import IntegrationHeartBeatFactory from apps.labels.tests.factories import ( AlertGroupAssociatedLabelFactory, @@ -107,8 +108,6 @@ from apps.webhooks.tests.test_webhook_presets import TEST_WEBHOOK_PRESET_ID, Tes register(OrganizationFactory) register(UserFactory) register(TeamFactory) - - register(AlertReceiveChannelFactory) register(AlertReceiveChannelConnectionFactory) register(ChannelFilterFactory) @@ -123,29 +122,24 @@ register(AlertGroupLogRecordFactory) register(InvitationFactory) register(CustomActionFactory) register(SlackUserGroupFactory) - register(SlackUserIdentityFactory) register(SlackTeamIdentityFactory) register(SlackMessageFactory) - register(TelegramToUserConnectorFactory) register(TelegramChannelFactory) register(TelegramVerificationCodeFactory) register(TelegramChannelVerificationCodeFactory) register(TelegramMessageFactory) - register(ResolutionNoteSlackMessageFactory) - register(PhoneCallRecordFactory) register(SMSRecordFactory) register(EmailMessageFactory) - register(IntegrationHeartBeatFactory) register(LiveSettingFactory) - register(LabelKeyFactory) register(LabelValueFactory) register(AlertReceiveChannelAssociatedLabelFactory) +register(GoogleOAuth2UserFactory) IS_RBAC_ENABLED = os.getenv("ONCALL_TESTING_RBAC_ENABLED", "True") == "True" @@ -1067,3 +1061,11 @@ def make_webhook_label_association(make_label_key_and_value): return WebhookAssociatedLabelFactory(webhook=webhook, organization=organization, key=key, value=value, **kwargs) return _make_integration_label_association + + +@pytest.fixture +def make_google_oauth2_user_for_user(): + def _make_google_oauth2_user_for_user(user): + return GoogleOAuth2UserFactory(user=user) + + return _make_google_oauth2_user_for_user