/
test_rate_limit.py
120 lines (103 loc) · 4.15 KB
/
test_rate_limit.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
"""Test for asyncprawcore.Sessions module."""
from copy import copy
from unittest.mock import patch
import pytest
from asyncprawcore.rate_limit import RateLimiter
from . import UnitTest
class TestRateLimiter(UnitTest):
@pytest.fixture
def rate_limiter(self):
rate_limiter = RateLimiter(window_size=600)
rate_limiter.next_request_timestamp = 100
return rate_limiter
@staticmethod
def _headers(remaining, used, reset):
return {
"x-ratelimit-remaining": str(float(remaining)),
"x-ratelimit-used": str(used),
"x-ratelimit-reset": str(reset),
}
@patch("time.time")
@patch("asyncio.sleep")
async def test_delay(self, mock_sleep, mock_time, rate_limiter):
mock_time.return_value = 1
await rate_limiter.delay()
assert mock_time.called
mock_sleep.assert_called_with(99)
@patch("time.time")
@patch("asyncio.sleep")
async def test_delay__no_sleep_when_time_in_past(
self, mock_sleep, mock_time, rate_limiter
):
mock_time.return_value = 101
await rate_limiter.delay()
assert mock_time.called
assert not mock_sleep.called
@patch("asyncio.sleep")
async def test_delay__no_sleep_when_time_is_not_set(self, mock_sleep, rate_limiter):
await rate_limiter.delay()
assert not mock_sleep.called
@patch("time.time")
@patch("asyncio.sleep")
async def test_delay__no_sleep_when_times_match(
self, mock_sleep, mock_time, rate_limiter
):
mock_time.return_value = 100
await rate_limiter.delay()
assert mock_time.called
assert not mock_sleep.called
@patch("time.time")
def test_update__compute_delay_with_no_previous_info(self, mock_time, rate_limiter):
mock_time.return_value = 100
rate_limiter.update(self._headers(60, 100, 60))
assert rate_limiter.remaining == 60
assert rate_limiter.used == 100
assert rate_limiter.next_request_timestamp == 100
@patch("time.time")
def test_update__compute_delay_with_single_client(self, mock_time, rate_limiter):
rate_limiter.remaining = 61
rate_limiter.window_size = 150
mock_time.return_value = 100
rate_limiter.update(self._headers(50, 100, 60))
assert rate_limiter.remaining == 50
assert rate_limiter.used == 100
assert rate_limiter.next_request_timestamp == 110
@patch("time.time")
def test_update__compute_delay_with_six_clients(self, mock_time, rate_limiter):
rate_limiter.remaining = 66
rate_limiter.window_size = 180
mock_time.return_value = 100
rate_limiter.update(self._headers(60, 100, 72))
assert rate_limiter.remaining == 60
assert rate_limiter.used == 100
assert rate_limiter.next_request_timestamp == 104.5
@patch("time.time")
def test_update__delay_full_time_with_negative_remaining(
self, mock_time, rate_limiter
):
mock_time.return_value = 37
rate_limiter.remaining = -1
rate_limiter.update(self._headers(0, 100, 13))
assert rate_limiter.remaining == 0
assert rate_limiter.used == 100
assert rate_limiter.next_request_timestamp == 50
@patch("time.time")
def test_update__delay_full_time_with_zero_remaining(self, mock_time, rate_limiter):
mock_time.return_value = 37
rate_limiter.remaining = 0
rate_limiter.update(self._headers(0, 100, 13))
assert rate_limiter.remaining == 0
assert rate_limiter.used == 100
assert rate_limiter.next_request_timestamp == 50
def test_update__no_change_without_headers(self, rate_limiter):
prev = copy(rate_limiter)
rate_limiter.update({})
assert prev.remaining == rate_limiter.remaining
assert prev.used == rate_limiter.used
assert rate_limiter.next_request_timestamp == prev.next_request_timestamp
def test_update__values_change_without_headers(self, rate_limiter):
rate_limiter.remaining = 10
rate_limiter.used = 99
rate_limiter.update({})
assert rate_limiter.remaining == 9
assert rate_limiter.used == 100