purldb/packagedb/tests/test_throttling.py at 8b8d061c5e060740143ad452f41904a7cc02cf5c · aboutcode-org/purldb · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
#
# Copyright (c) nexB Inc. and others. All rights reserved.
# PurlDB is a trademark of nexB Inc.
# SPDX-License-Identifier: Apache-2.0
# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
# See https://github.com/aboutcode-org/purldb for support or download.
# See https://aboutcode.org for more information about nexB OSS projects.
#

from unittest.mock import patch

from django.contrib.auth.models import User
from django.core.cache import cache
from rest_framework.test import APIClient
from rest_framework.test import APITestCase


@patch("rest_framework.throttling.UserRateThrottle.get_rate", lambda x: "20/day")
@patch("rest_framework.throttling.AnonRateThrottle.get_rate", lambda x: "10/day")
class ThrottleApiTests(APITestCase):

    def setUp(self):
        cache.clear()
        # create a basic user
        self.user = User.objects.create_user(
            username="username",
            email="e@mail.com",
            password="secret",  # NOQA
        )
        self.auth = f"Token {self.user.auth_token.key}"
        self.csrf_client = APIClient(enforce_csrf_checks=True)
        self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)

        # create a staff user
        self.staff_user = User.objects.create_user(
            username="staff_username",
            email="staff_e@mail.com",
            password="secret",  # NOQA
            is_staff=True,
        )
        self.staff_auth = f"Token {self.staff_user.auth_token.key}"
        self.staff_csrf_client = APIClient(enforce_csrf_checks=True)
        self.staff_csrf_client.credentials(HTTP_AUTHORIZATION=self.staff_auth)

        self.csrf_client_anon = APIClient(enforce_csrf_checks=True)

    def test_package_endpoint_throttling(self):
        for i in range(0, 20):
            response = self.csrf_client.get("/api/packages/")
            self.assertEqual(response.status_code, 200)
            response = self.staff_csrf_client.get("/api/packages/")
            self.assertEqual(response.status_code, 200)

        response = self.csrf_client.get("/api/packages/")
        # 429 - too many requests for basic user
        self.assertEqual(response.status_code, 429)

        response = self.staff_csrf_client.get("/api/packages/", format="json")
        # 200 - staff user can access API unlimited times
        self.assertEqual(response.status_code, 200)

        # A anonymous user can only access /packages endpoint 10 times a day
        for i in range(0, 10):
            response = self.csrf_client_anon.get("/api/packages/")
            self.assertEqual(response.status_code, 200)

        response = self.csrf_client_anon.get("/api/packages/")
        # 429 - too many requests for anon user
        self.assertEqual(response.status_code, 429)
        self.assertEqual(
            response.data.get("message"),
            "Your request has been throttled. Please contact support@nexb.com",
        )

        response = self.csrf_client_anon.get("/api/resources/")
        # 429 - too many requests for anon user
        self.assertEqual(response.status_code, 429)
        self.assertEqual(
            response.data.get("message"),
            "Your request has been throttled. Please contact support@nexb.com",
        )