Files
zulip/zerver/tests/test_external.py
Mateusz Mandera 4e9f77a6c4 rate_limit: Adjust keys() of some RateLimitedObjects.
type().__name__ is sufficient, and much readable than type(), so it's
better to use the former for keys.
We also make the classes consistent in forming the keys in the format
type(self).__name__:identifier and adjust logger.warning and statsd to
take advantage of that and simply log the key().
2020-03-22 18:42:35 -07:00

126 lines
5.4 KiB
Python

# -*- coding: utf-8 -*-
from django.conf import settings
from django.core.exceptions import ValidationError
from django.http import HttpResponse
from zerver.forms import email_is_not_mit_mailing_list
from zerver.lib.rate_limiter import (
add_ratelimit_rule,
remove_ratelimit_rule,
RateLimitedUser,
RateLimiterLockingException,
)
from zerver.lib.zephyr import compute_mit_user_fullname
from zerver.lib.test_classes import (
ZulipTestCase,
)
from zerver.models import (
UserProfile,
)
import DNS
import mock
import time
class MITNameTest(ZulipTestCase):
def test_valid_hesiod(self) -> None:
with mock.patch('DNS.dnslookup', return_value=[['starnine:*:84233:101:Athena Consulting Exchange User,,,:/mit/starnine:/bin/bash']]):
self.assertEqual(compute_mit_user_fullname(self.mit_email("starnine")), "Athena Consulting Exchange User")
with mock.patch('DNS.dnslookup', return_value=[['sipbexch:*:87824:101:Exch Sipb,,,:/mit/sipbexch:/bin/athena/bash']]):
self.assertEqual(compute_mit_user_fullname("sipbexch@mit.edu"), "Exch Sipb")
def test_invalid_hesiod(self) -> None:
with mock.patch('DNS.dnslookup', side_effect=DNS.Base.ServerError('DNS query status: NXDOMAIN', 3)):
self.assertEqual(compute_mit_user_fullname("1234567890@mit.edu"), "1234567890@mit.edu")
with mock.patch('DNS.dnslookup', side_effect=DNS.Base.ServerError('DNS query status: NXDOMAIN', 3)):
self.assertEqual(compute_mit_user_fullname("ec-discuss@mit.edu"), "ec-discuss@mit.edu")
def test_mailinglist(self) -> None:
with mock.patch('DNS.dnslookup', side_effect=DNS.Base.ServerError('DNS query status: NXDOMAIN', 3)):
self.assertRaises(ValidationError, email_is_not_mit_mailing_list, "1234567890@mit.edu")
with mock.patch('DNS.dnslookup', side_effect=DNS.Base.ServerError('DNS query status: NXDOMAIN', 3)):
self.assertRaises(ValidationError, email_is_not_mit_mailing_list, "ec-discuss@mit.edu")
def test_notmailinglist(self) -> None:
with mock.patch('DNS.dnslookup', return_value=[['POP IMAP.EXCHANGE.MIT.EDU starnine']]):
email_is_not_mit_mailing_list("sipbexch@mit.edu")
class RateLimitTests(ZulipTestCase):
def setUp(self) -> None:
super().setUp()
settings.RATE_LIMITING = True
add_ratelimit_rule(1, 5)
def tearDown(self) -> None:
settings.RATE_LIMITING = False
remove_ratelimit_rule(1, 5)
super().tearDown()
def send_api_message(self, user: UserProfile, content: str) -> HttpResponse:
return self.api_post(user, "/api/v1/messages", {"type": "stream",
"to": "Verona",
"client": "test suite",
"content": content,
"topic": "whatever"})
def test_headers(self) -> None:
user = self.example_user('hamlet')
RateLimitedUser(user).clear_history()
result = self.send_api_message(user, "some stuff")
self.assertTrue('X-RateLimit-Remaining' in result)
self.assertTrue('X-RateLimit-Limit' in result)
self.assertTrue('X-RateLimit-Reset' in result)
def test_ratelimit_decrease(self) -> None:
user = self.example_user('hamlet')
RateLimitedUser(user).clear_history()
result = self.send_api_message(user, "some stuff")
limit = int(result['X-RateLimit-Remaining'])
result = self.send_api_message(user, "some stuff 2")
newlimit = int(result['X-RateLimit-Remaining'])
self.assertEqual(limit, newlimit + 1)
def test_hit_ratelimits(self) -> None:
user = self.example_user('cordelia')
RateLimitedUser(user).clear_history()
start_time = time.time()
for i in range(6):
with mock.patch('time.time', return_value=(start_time + i * 0.1)):
result = self.send_api_message(user, "some stuff %s" % (i,))
self.assertEqual(result.status_code, 429)
json = result.json()
self.assertEqual(json.get("result"), "error")
self.assertIn("API usage exceeded rate limit", json.get("msg"))
self.assertEqual(json.get('retry-after'), 0.5)
self.assertTrue('Retry-After' in result)
self.assertEqual(result['Retry-After'], '0.5')
# We actually wait a second here, rather than force-clearing our history,
# to make sure the rate-limiting code automatically forgives a user
# after some time has passed.
with mock.patch('time.time', return_value=(start_time + 1.0)):
result = self.send_api_message(user, "Good message")
self.assert_json_success(result)
@mock.patch('zerver.lib.rate_limiter.logger.warning')
def test_hit_ratelimiterlockingexception(self, mock_warn: mock.MagicMock) -> None:
user = self.example_user('cordelia')
RateLimitedUser(user).clear_history()
with mock.patch('zerver.lib.rate_limiter.RedisRateLimiterBackend.incr_ratelimit',
side_effect=RateLimiterLockingException):
result = self.send_api_message(user, "some stuff")
self.assertEqual(result.status_code, 429)
mock_warn.assert_called_with("Deadlock trying to incr_ratelimit for RateLimitedUser:%s:api_by_user"
% (user.id,))