|
15 | 15 | import time |
16 | 16 | import threading |
17 | 17 |
|
18 | | -from typing import Callable, Optional |
| 18 | +from typing import Callable |
19 | 19 | from functools import wraps |
20 | 20 | from collections import deque |
21 | 21 | from enum import Enum |
@@ -105,22 +105,91 @@ def run_when_exceed_timeout(wait: int, until: int): |
105 | 105 | print("threading id: {} I am a long time hook".format(threading.get_ident())) |
106 | 106 |
|
107 | 107 |
|
108 | | -@RateLimiter(max_call=2, period=1, exceed_strategy=EStrategy.Sleep, exceed_hook=run_when_exceed) |
109 | | -def work(*args, **kwargs): |
110 | | - print(args, kwargs) |
111 | | - |
112 | | - |
113 | 108 | def work_timeout(*args, **kwargs): |
114 | 109 | print(args, kwargs) |
115 | 110 | time.sleep(1) |
116 | 111 |
|
| 112 | +# ==================================================== |
| 113 | +# TEST CASE |
| 114 | +# ==================================================== |
| 115 | +import pytest |
117 | 116 |
|
118 | | -if __name__ == "__main__": |
119 | 117 |
|
120 | | - for i in range(30): |
121 | | - t = threading.Thread(target=work, args=(1,2), kwargs={"a": 3}) |
122 | | - t.setDaemon(True) |
123 | | - t.start() |
| 118 | +def test_simple_raise(): |
| 119 | + print("\n>>>>>>> log for test_simple_raise") |
| 120 | + |
| 121 | + result = [] |
| 122 | + limiter = RateLimiter(max_call=2, period=1, exceed_strategy=EStrategy.Raise) |
| 123 | + |
| 124 | + def worker(*args, **kwargs): |
| 125 | + print("I am doing something at ", time.time()) |
| 126 | + |
| 127 | + with pytest.raises(Exception): |
| 128 | + for i in range(3): |
| 129 | + with limiter: |
| 130 | + result.append(worker()) |
| 131 | + |
| 132 | + assert len(result) == 2 |
| 133 | + |
| 134 | + |
| 135 | +def test_simple_sleep(): |
| 136 | + print("\n>>>>>>> log for test_simple_sleep") |
| 137 | + |
| 138 | + result = [] |
| 139 | + |
| 140 | + @RateLimiter(max_call=2, period=1, exceed_strategy=EStrategy.Sleep) |
| 141 | + def work(*args, **kwargs): |
| 142 | + print("I am doing something at ", time.time()) |
| 143 | + |
| 144 | + def check_after_5_seconds(): |
| 145 | + time.sleep(5) |
| 146 | + assert len(result) == 10 |
| 147 | + time.sleep(3600) |
| 148 | + |
| 149 | + t = threading.Thread(target=check_after_5_seconds) |
| 150 | + t.setDaemon(True) |
| 151 | + t.start() |
| 152 | + |
| 153 | + for i in range(11): |
| 154 | + result.append(work()) |
| 155 | + assert t.is_alive() |
| 156 | + |
| 157 | + assert len(result) == 11 |
| 158 | + |
| 159 | + |
| 160 | +def test_thread_safe(): |
| 161 | + print("\n>>>>>>> log for test_thread_safe") |
| 162 | + |
| 163 | + result = [] |
| 164 | + |
| 165 | + @RateLimiter(max_call=2, period=1, exceed_strategy=EStrategy.Sleep) |
| 166 | + def work(*args, **kwargs): |
| 167 | + print("I am doing something at ", time.time()) |
| 168 | + |
| 169 | + def check_after_5_seconds(): |
| 170 | + time.sleep(5) |
| 171 | + assert len(result) == 10 |
| 172 | + time.sleep(3600) |
| 173 | + |
| 174 | + def thread_work(): |
| 175 | + for _ in range(2): |
| 176 | + result.append(work()) |
| 177 | + |
| 178 | + t = threading.Thread(target=check_after_5_seconds) |
| 179 | + t.setDaemon(True) |
| 180 | + t.start() |
| 181 | + |
| 182 | + w = threading.Thread(target=thread_work) |
| 183 | + w.setDaemon(True) |
| 184 | + w.start() |
| 185 | + |
| 186 | + for _ in range(11): |
| 187 | + result.append(work()) |
| 188 | + assert t.is_alive() |
| 189 | + |
| 190 | + assert len(result) == 13 |
| 191 | + |
| 192 | + |
| 193 | +if __name__ == "__main__": |
| 194 | + pytest.main(["./ratelimiter.py", "-v", "-s"]) |
124 | 195 |
|
125 | | - for i in range(100): |
126 | | - work(1, 2, a=3) |
|
0 commit comments