Add tests

engine/apps/user_management/middlewares.py

@@ -4,7 +4,7 @@ import re
 import requests
 from django.http import HttpResponse
 from django.utils.deprecation import MiddlewareMixin
-from rest_framework.status import HTTP_500_INTERNAL_SERVER_ERROR
+from rest_framework import status
 
 from apps.user_management.models.region import OrganizationMovedException
 from common.api_helpers.utils import create_engine_url
@@ -18,7 +18,7 @@ class OrganizationMovedMiddleware(MiddlewareMixin):
             region = exception.organization.migration_destination
             if not region.oncall_backend_url:
                 return HttpResponse(
-                    "Organization migration destination undefined URL", status=HTTP_500_INTERNAL_SERVER_ERROR
+                    "Organization migration destination undefined URL", status=status.HTTP_500_INTERNAL_SERVER_ERROR
                 )
 
             url = create_engine_url(request.path, override_base=region.oncall_backend_url)
@@ -30,15 +30,17 @@ class OrganizationMovedMiddleware(MiddlewareMixin):
                 (regex.sub("", header), value) for (header, value) in request.META.items() if header.startswith("HTTP_")
             )
 
-            if request.method == "GET":
-                response = requests.get(url, headers=headers)
-            elif request.method == "POST":
-                response = requests.post(url, data=request.body, headers=headers)
-            elif request.method == "PUT":
-                response = requests.put(url, data=request.body, headers=headers)
-            elif request.method == "DELETE":
-                response = requests.delete(url, headers=headers)
-            elif request.method == "OPTIONS":
-                response = requests.options(url, headers=headers)
-
+            response = self.make_request(request.method, url, headers, request.body)
             return HttpResponse(response.content, status=response.status_code)
+
+    def make_request(self, method, url, headers, body):
+        if method == "GET":
+            return requests.get(url, headers=headers)
+        elif method == "POST":
+            return requests.post(url, data=body, headers=headers)
+        elif method == "PUT":
+            return requests.put(url, data=body, headers=headers)
+        elif method == "DELETE":
+            return requests.delete(url, headers=headers)
+        elif method == "OPTIONS":
+            return requests.options(url, headers=headers)

engine/apps/user_management/tests/factories.py

@@ -1,6 +1,6 @@
 import factory
 
-from apps.user_management.models import Organization, Team, User
+from apps.user_management.models import Organization, Region, Team, User
 from common.utils import UniqueFaker
 
 
@@ -31,3 +31,12 @@ class TeamFactory(factory.DjangoModelFactory):
 
     class Meta:
         model = Team
+
+
+class RegionFactory(factory.DjangoModelFactory):
+    name = factory.Faker("country")
+    slug = factory.Faker("slug")
+    oncall_backend_url = factory.Faker("url")
+
+    class Meta:
+        model = Region

engine/apps/user_management/tests/test_region.py

@@ -0,0 +1,219 @@
+from unittest.mock import patch
+
+import pytest
+from django.http import HttpResponse
+from django.urls import reverse
+from rest_framework import status
+from rest_framework.test import APIClient
+
+from apps.alerts.models import AlertReceiveChannel
+from apps.auth_token.auth import ApiTokenAuthentication, ScheduleExportAuthentication, UserScheduleExportAuthentication
+from apps.auth_token.models import ScheduleExportAuthToken, UserScheduleExportAuthToken
+from apps.integrations.views import AlertManagerAPIView
+from apps.schedules.models import OnCallScheduleWeb
+from apps.user_management.models.region import OrganizationMovedException
+
+
+@pytest.mark.django_db
+def test_organization_region_delete(
+    make_organization_and_region,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    organization.refresh_from_db()
+    assert organization.migration_destination.slug == region.slug
+    region.delete()
+
+    organization.refresh_from_db()
+    assert organization.migration_destination is None
+
+
+@pytest.mark.django_db
+def test_integration_does_not_raise_exception_organization_moved(
+    make_organization,
+    make_alert_receive_channel,
+):
+    organization = make_organization()
+    alert_receive_channel = make_alert_receive_channel(
+        organization=organization,
+        integration=AlertReceiveChannel.INTEGRATION_ALERTMANAGER,
+    )
+
+    try:
+        am = AlertManagerAPIView()
+        am.dispatch(alert_channel_key=alert_receive_channel.token)
+        assert False
+    except OrganizationMovedException:
+        assert False
+    except Exception:
+        assert True
+
+
+@pytest.mark.django_db
+def test_integration_raises_exception_organization_moved(
+    make_organization_and_region,
+    make_alert_receive_channel,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    alert_receive_channel = make_alert_receive_channel(
+        organization=organization,
+        integration=AlertReceiveChannel.INTEGRATION_ALERTMANAGER,
+    )
+
+    try:
+        am = AlertManagerAPIView()
+        am.dispatch(alert_channel_key=alert_receive_channel.token)
+        assert False
+    except OrganizationMovedException as e:
+        assert e.organization == organization
+
+
+@patch("apps.user_management.middlewares.OrganizationMovedMiddleware.make_request")
+@pytest.mark.django_db
+def test_organization_moved_middleware(
+    mocked_make_request,
+    make_organization_and_region,
+    make_alert_receive_channel,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    alert_receive_channel = make_alert_receive_channel(
+        organization=organization,
+        integration=AlertReceiveChannel.INTEGRATION_ALERTMANAGER,
+    )
+
+    expected_message = bytes(f"Redirected to {region.oncall_backend_url}", 'utf-8')
+    mocked_make_request.return_value = HttpResponse(expected_message, status=status.HTTP_200_OK)
+
+    client = APIClient()
+    url = reverse("integrations:alertmanager", kwargs={"alert_channel_key": alert_receive_channel.token})
+
+    data = {"value": "test"}
+    response = client.post(url, data, format="json")
+    assert mocked_make_request.called
+    assert response.content == expected_message
+    assert response.status_code == status.HTTP_200_OK
+
+
+@pytest.mark.django_db
+def test_api_token_does_not_raise_exception_organization_moved(
+    make_organization,
+    make_user_for_organization,
+    make_public_api_token,
+):
+    organization = make_organization()
+
+    admin = make_user_for_organization(organization)
+    _, token = make_public_api_token(admin, organization)
+
+    try:
+        api_auth = ApiTokenAuthentication()
+        api_auth.authenticate_credentials(token)
+        assert True
+    except OrganizationMovedException:
+        assert False
+
+
+@pytest.mark.django_db
+def test_api_token_raises_exception_organization_moved(
+    make_organization_and_region,
+    make_user_for_organization,
+    make_public_api_token,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    admin = make_user_for_organization(organization)
+    _, token = make_public_api_token(admin, organization)
+
+    try:
+        api_auth = ApiTokenAuthentication()
+        api_auth.authenticate_credentials(token)
+        assert False
+    except OrganizationMovedException as e:
+        assert e.organization == organization
+
+
+@pytest.mark.django_db
+def test_schedule_export_token_does_not_raise_exception_organization_moved(
+    make_organization,
+    make_user_for_organization,
+    make_public_api_token,
+    make_schedule,
+):
+    organization = make_organization()
+    schedule = make_schedule(organization, schedule_class=OnCallScheduleWeb)
+
+    admin = make_user_for_organization(organization)
+    _, token = ScheduleExportAuthToken.create_auth_token(admin, organization, schedule)
+
+    try:
+        schedule_auth = ScheduleExportAuthentication()
+        schedule_auth.authenticate_credentials(token, schedule.public_primary_key)
+        assert True
+    except OrganizationMovedException:
+        assert False
+
+
+@pytest.mark.django_db
+def test_schedule_export_token_raises_exception_organization_moved(
+    make_organization_and_region,
+    make_user_for_organization,
+    make_public_api_token,
+    make_schedule,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+    schedule = make_schedule(organization, schedule_class=OnCallScheduleWeb)
+
+    admin = make_user_for_organization(organization)
+    _, token = ScheduleExportAuthToken.create_auth_token(admin, organization, schedule)
+
+    try:
+        schedule_auth = ScheduleExportAuthentication()
+        schedule_auth.authenticate_credentials(token, schedule.public_primary_key)
+        assert False
+    except OrganizationMovedException as e:
+        assert e.organization == organization
+
+
+@pytest.mark.django_db
+def test_user_schedule_export_token_does_not_raise_exception_organization_moved(
+    make_organization,
+    make_user_for_organization,
+    make_public_api_token,
+):
+    organization = make_organization()
+    admin = make_user_for_organization(organization)
+    _, token = UserScheduleExportAuthToken.create_auth_token(admin, organization)
+
+    try:
+        user_schedule_auth = UserScheduleExportAuthentication()
+        user_schedule_auth.authenticate_credentials(token, admin.public_primary_key)
+        assert True
+    except OrganizationMovedException:
+        assert False
+
+
+@pytest.mark.django_db
+def test_user_schedule_export_token_raises_exception_organization_moved(
+    make_organization_and_region,
+    make_user_for_organization,
+    make_public_api_token,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    admin = make_user_for_organization(organization)
+    _, token = UserScheduleExportAuthToken.create_auth_token(admin, organization)
+
+    try:
+        user_schedule_auth = UserScheduleExportAuthentication()
+        user_schedule_auth.authenticate_credentials(token, admin.public_primary_key)
+        assert False
+    except OrganizationMovedException as e:
+        assert e.organization == organization

engine/conftest.py

@@ -70,7 +70,7 @@ from apps.telegram.tests.factories import (
 )
 from apps.twilioapp.tests.factories import PhoneCallFactory, SMSFactory
 from apps.user_management.models.user import User, listen_for_user_model_save
-from apps.user_management.tests.factories import OrganizationFactory, TeamFactory, UserFactory
+from apps.user_management.tests.factories import OrganizationFactory, RegionFactory, TeamFactory, UserFactory
 from common.constants.role import Role
 
 register(OrganizationFactory)
@@ -666,3 +666,23 @@ def load_slack_urls(settings):
         reload(sys.modules[urlconf])
     else:
         import_module(urlconf)
+
+
+@pytest.fixture
+def make_region():
+    def _make_region(**kwargs):
+        region = RegionFactory(**kwargs)
+        return region
+
+    return _make_region
+
+
+@pytest.fixture
+def make_organization_and_region(make_organization, make_region):
+    def _make_organization_and_region():
+        organization = make_organization()
+        region = make_region()
+        organization.migration_destination = region
+        return organization, region
+
+    return _make_organization_and_region