-
-
Notifications
You must be signed in to change notification settings - Fork 331
/
tests_utils_ratelimit.py
151 lines (116 loc) · 4.57 KB
/
tests_utils_ratelimit.py
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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import hashlib
from django.core.cache import cache
from django.test import TestCase, RequestFactory
from django.contrib.auth.models import AnonymousUser, User
from django.contrib.messages.storage.fallback import FallbackStorage
from django.conf import settings
from django.core.cache import caches
from ..utils.ratelimit import RateLimit
from ..utils.ratelimit.decorators import ratelimit
def setup_request_factory_messages(req):
# RequestFactory: It does not support middleware.
# Session and authentication attributes must be supplied
# by the test itself if required for the view to function properly.
setattr(req, 'session', 'session')
messages = FallbackStorage(req)
setattr(req, '_messages', messages)
class UtilsRateLimitTests(TestCase):
def setUp(self):
cache.clear()
def test_rate_limit_split_rate(self):
req = RequestFactory().post('/')
req.user = AnonymousUser()
rl = RateLimit(req, 'func_name')
self.assertEqual(rl.split_rate('5/m'), (5, 60))
self.assertEqual(rl.split_rate('5/5m'), (5, 60 * 5))
self.assertEqual(rl.split_rate('5/s'), (5, 1))
self.assertEqual(rl.split_rate('5/5s'), (5, 1 * 5))
self.assertEqual(rl.split_rate('5/15s'), (5, 15))
self.assertEqual(rl.split_rate('15/15s'), (15, 15))
def test_rate_limit_user_or_ip(self):
req = RequestFactory().get('/')
req.user = AnonymousUser()
setup_request_factory_messages(req)
@ratelimit(method=['GET', ], rate='1/m')
def limit_ip(request):
return request.is_limited
self.assertFalse(limit_ip(req))
self.assertTrue(limit_ip(req))
req.user = User()
req.user.pk = 1
self.assertFalse(limit_ip(req))
self.assertTrue(limit_ip(req))
def test_rate_limit_method(self):
rf = RequestFactory()
post = rf.post('/')
get = rf.get('/')
post.user = AnonymousUser()
get.user = AnonymousUser()
setup_request_factory_messages(get)
setup_request_factory_messages(post)
@ratelimit(method=['POST', ], rate='1/m')
def limit_post(request):
return request.is_limited
self.assertFalse(limit_post(post))
self.assertTrue(limit_post(post))
self.assertFalse(limit_post(get))
def test_rate_limit_field(self):
req = RequestFactory().post('/', {'username': 'esteban', })
req.user = AnonymousUser()
setup_request_factory_messages(req)
@ratelimit(field='username', rate='1/m')
def username(request):
return request.is_limited
self.assertFalse(username(req))
self.assertTrue(username(req))
req.user = User()
req.user.pk = 1
self.assertTrue(username(req))
def test_rate_limit_field_empty(self):
empty = RequestFactory().post('/', {'username': '', })
empty.user = AnonymousUser()
@ratelimit(field='username', rate='1/m')
def username(request):
return request.is_limited
self.assertFalse(username(empty))
empty.user = User()
empty.user.pk = 1
self.assertFalse(username(empty))
def test_rate_limit_rate(self):
req = RequestFactory().post('/')
req.user = AnonymousUser()
setup_request_factory_messages(req)
@ratelimit(rate='2/m')
def two(request):
return request.is_limited
self.assertFalse(two(req))
self.assertFalse(two(req))
self.assertTrue(two(req))
def test_rate_limit_hash_key(self):
"""
Keys should be stored as hashes
"""
req = RequestFactory().post('/')
req.user = User()
req.user.pk = 1
RateLimit(req, 'func_name')
rl_cache = caches[settings.ST_RATELIMIT_CACHE]
self.assertIsNotNone(rl_cache.get('srl:02b3cee0bd2a40ec0fca9b1bef06fb560a081673'))
def test_rate_limit_unique_key(self):
"""
Keys should contain the full module path and function name
"""
req = RequestFactory().post('/')
req.user = User()
req.user.pk = 1
@ratelimit(rate='1/m')
def one(request):
pass
key_part = '%s.%s:user:%d' % (one.__module__, one.__name__, req.user.pk)
key_hash = hashlib.sha1(key_part.encode('utf-8')).hexdigest()
key = '%s:%s' % (settings.ST_RATELIMIT_CACHE_PREFIX, key_hash)
one(req)
rl_cache = caches[settings.ST_RATELIMIT_CACHE]
self.assertIsNotNone(rl_cache.get(key))