zulip/zerver/tests/test_external.py

# -*- coding: utf-8 -*-
from __future__ import absolute_import
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,
    clear_user_history,
    remove_ratelimit_rule,
)

from zerver.lib.actions import compute_mit_user_fullname
from zerver.lib.test_classes import (
    ZulipTestCase,
)

import DNS
import mock
import time
import ujson

from six.moves import urllib
from six.moves import range
from typing import Text

class MITNameTest(ZulipTestCase):
    def test_valid_hesiod(self):
        # type: () -> 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("starnine@mit.edu"), "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):
        # type: () -> 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):
        # type: () -> 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):
        # type: () -> 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):
        # type: () -> None
        settings.RATE_LIMITING = True
        add_ratelimit_rule(1, 5)

    def tearDown(self):
        # type: () -> None
        settings.RATE_LIMITING = False
        remove_ratelimit_rule(1, 5)

    def send_api_message(self, email, content):
        # type: (Text, Text) -> HttpResponse
        return self.client_post("/api/v1/messages", {"type": "stream",
                                                     "to": "Verona",
                                                     "client": "test suite",
                                                     "content": content,
                                                     "subject": "Test subject"},
                                **self.api_auth(email))

    def test_headers(self):
        # type: () -> None
        user = self.example_user('hamlet')
        email = user.email
        clear_user_history(user)

        result = self.send_api_message(email, "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):
        # type: () -> None
        user = self.example_user('hamlet')
        email = user.email
        clear_user_history(user)
        result = self.send_api_message(email, "some stuff")
        limit = int(result['X-RateLimit-Remaining'])

        result = self.send_api_message(email, "some stuff 2")
        newlimit = int(result['X-RateLimit-Remaining'])
        self.assertEqual(limit, newlimit + 1)

    def test_hit_ratelimits(self):
        # type: () -> None
        user = self.example_user('cordelia')
        email = user.email
        clear_user_history(user)

        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(email, "some stuff %s" % (i,))

        self.assertEqual(result.status_code, 429)
        json = ujson.loads(result.content)
        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(email, "Good message")

            self.assert_json_success(result)