oncall-engine/engine/apps/integrations/tests/test_ratelimit.py
Michael Derynck 6b40f95033 World, meet OnCall!
Co-authored-by: Eve832 <eve.meelan@grafana.com>
    Co-authored-by: Francisco Montes de Oca <nevermind89x@gmail.com>
    Co-authored-by: Ildar Iskhakov <ildar.iskhakov@grafana.com>
    Co-authored-by: Innokentii Konstantinov <innokenty.konstantinov@grafana.com>
    Co-authored-by: Julia <ferril.darkdiver@gmail.com>
    Co-authored-by: maskin25 <kengurek@gmail.com>
    Co-authored-by: Matias Bordese <mbordese@gmail.com>
    Co-authored-by: Matvey Kukuy <motakuk@gmail.com>
    Co-authored-by: Michael Derynck <michael.derynck@grafana.com>
    Co-authored-by: Richard Hartmann <richih@richih.org>
    Co-authored-by: Robby Milo <robbymilo@fastmail.com>
    Co-authored-by: Timur Olzhabayev <timur.olzhabayev@grafana.com>
    Co-authored-by: Vadim Stepanov <vadimkerr@gmail.com>
    Co-authored-by: Yulia Shanyrova <yulia.shanyrova@grafana.com>
2022-06-03 08:09:47 -06:00

99 lines
3.2 KiB
Python

from unittest import mock
import pytest
from django.core.cache import cache
from django.test import Client
from django.urls import reverse
from apps.alerts.models import AlertReceiveChannel
# Ratelimit keys are stored in cache. Clean it before and after every test to make them idempotent.
def setup_module(module):
cache.clear()
def teardown_module(module):
cache.clear()
@mock.patch("ratelimit.utils._split_rate", return_value=(1, 60))
@mock.patch("apps.integrations.tasks.create_alert.apply_async", return_value=None)
@pytest.mark.django_db
def test_ratelimit_alerts_per_integration(
mocked_task,
mocked_rate,
make_organization,
make_alert_receive_channel,
):
organization = make_organization()
integration = make_alert_receive_channel(organization, integration=AlertReceiveChannel.INTEGRATION_WEBHOOK)
url = reverse(
"integrations:universal",
kwargs={"integration_type": AlertReceiveChannel.INTEGRATION_WEBHOOK, "alert_channel_key": integration.token},
)
c = Client()
response = c.post(url, data={"message": "This is the test alert from amixr"})
assert response.status_code == 200
response = c.post(url, data={"message": "This is the test alert from amixr"})
assert response.status_code == 429
assert mocked_task.call_count == 1
@pytest.mark.skip(reason="SQLITE Incompatibility")
@mock.patch("ratelimit.utils._split_rate", return_value=(1, 60))
@mock.patch("apps.integrations.tasks.create_alert.apply_async", return_value=None)
@pytest.mark.django_db
def test_ratelimit_alerts_per_team(
mocked_task,
mocked_rate,
make_organization,
make_alert_receive_channel,
):
organization = make_organization()
integration_1 = make_alert_receive_channel(organization, integration=AlertReceiveChannel.INTEGRATION_WEBHOOK)
url_1 = reverse("integrations:webhook", kwargs={"alert_channel_key": integration_1.token})
integration_2 = make_alert_receive_channel(organization, integration=AlertReceiveChannel.INTEGRATION_WEBHOOK)
url_2 = reverse("integrations:webhook", kwargs={"alert_channel_key": integration_2.token})
c = Client()
response = c.post(url_1, data={"message": "This is the test alert from amixr"})
assert response.status_code == 200
response = c.post(url_2, data={"message": "This is the test alert from amixr"})
assert response.status_code == 429
assert mocked_task.call_count == 1
@pytest.mark.skip(reason="SQLITE Incompatibility")
@mock.patch("ratelimit.utils._split_rate", return_value=(1, 60))
@mock.patch("apps.heartbeat.tasks.process_heartbeat_task.apply_async", return_value=None)
@pytest.mark.django_db
def test_ratelimit_integration_heartbeats(
mocked_task,
mocked_rate,
make_organization,
make_alert_receive_channel,
):
organization = make_organization()
integration = make_alert_receive_channel(organization, integration=AlertReceiveChannel.INTEGRATION_WEBHOOK)
url = reverse("integrations:webhook_heartbeat", kwargs={"alert_channel_key": integration.token})
c = Client()
response = c.post(url)
assert response.status_code == 200
response = c.post(url)
assert response.status_code == 429
response = c.get(url)
assert response.status_code == 429
assert mocked_task.call_count == 1