Coverage for agentos/tests/test_rate_limiter.py: 0%

1"""Comprehensive tests for agentos/tools/rate_limiter.py.""" 

2 

3import time 

4 

5import pytest 

6 

7from agentos.tools.rate_limiter import ( 

8 RateLimitExceeded, 

9 SlidingWindow, 

10 TokenBucket, 

11) 

12 

13# ============================================================================ 

14# RateLimitExceeded 

15# ============================================================================ 

16 

17 

18class TestRateLimitExceeded: 

19 def test_constructor_and_attributes(self): 

20 exc = RateLimitExceeded(key="api:test", limit=10.0, window=60.0) 

21 assert exc.key == "api:test" 

22 assert exc.limit == 10.0 

23 assert exc.window == 60.0 

24 

25 def test_string_message(self): 

26 exc = RateLimitExceeded(key="user:42", limit=5.0, window=1.0) 

27 msg = str(exc) 

28 assert "user:42" in msg 

29 assert "5.0" in msg 

30 assert "1.0" in msg 

31 

32 

33# ============================================================================ 

34# TokenBucket 

35# ============================================================================ 

36 

37 

38class TestTokenBucketInit: 

39 def test_default_burst_equals_rate(self): 

40 tb = TokenBucket(rate=10.0) 

41 assert tb._rate == 10.0 

42 assert tb._burst == 10.0 

43 

44 def test_custom_burst(self): 

45 tb = TokenBucket(rate=5.0, burst=20.0) 

46 assert tb._rate == 5.0 

47 assert tb._burst == 20.0 

48 

49 def test_rate_zero_raises(self): 

50 with pytest.raises(ValueError, match="positive"): 

51 TokenBucket(rate=0) 

52 

53 def test_rate_negative_raises(self): 

54 with pytest.raises(ValueError, match="positive"): 

55 TokenBucket(rate=-1.0) 

56 

57 def test_rate_property(self): 

58 tb = TokenBucket(rate=7.5) 

59 assert tb.rate == 7.5 

60 

61 

62class TestTokenBucketTryAcquire: 

63 def test_first_acquire_succeeds(self): 

64 tb = TokenBucket(rate=10.0, burst=5.0) 

65 assert tb.try_acquire("key1") is True 

66 

67 def test_acquire_within_burst(self): 

68 tb = TokenBucket(rate=10.0, burst=5.0) 

69 for _ in range(5): 

70 assert tb.try_acquire("key1") is True 

71 # 6th should fail 

72 assert tb.try_acquire("key1") is False 

73 

74 def test_acquire_multi_tokens(self): 

75 tb = TokenBucket(rate=10.0, burst=10.0) 

76 assert tb.try_acquire("key1", tokens=5.0) is True 

77 assert tb.try_acquire("key1", tokens=5.0) is True 

78 assert tb.try_acquire("key1", tokens=5.0) is False 

79 

80 def test_different_keys_independent(self): 

81 tb = TokenBucket(rate=10.0, burst=2.0) 

82 assert tb.try_acquire("key1") is True 

83 assert tb.try_acquire("key1") is True 

84 assert tb.try_acquire("key1") is False 

85 # key2 should be unaffected 

86 assert tb.try_acquire("key2") is True 

87 

88 def test_refill_after_wait(self): 

89 tb = TokenBucket(rate=100.0, burst=2.0) # 100 tokens/sec 

90 assert tb.try_acquire("key1") is True 

91 assert tb.try_acquire("key1") is True 

92 assert tb.try_acquire("key1") is False 

93 time.sleep(0.02) # ~2 tokens refilled 

94 assert tb.try_acquire("key1") is True 

95 

96 def test_refill_capped_at_burst(self): 

97 tb = TokenBucket(rate=1000.0, burst=3.0) 

98 for _ in range(3): 

99 assert tb.try_acquire("key1") is True 

100 assert tb.try_acquire("key1") is False 

101 time.sleep(0.1) # would refill 100, but burst caps at 3 

102 # should have exactly 3 tokens, not 100 

103 for _ in range(3): 

104 assert tb.try_acquire("key1") is True 

105 assert tb.try_acquire("key1") is False 

106 

107 def test_acquire_exact_tokens(self): 

108 tb = TokenBucket(rate=10.0, burst=10.0) 

109 assert tb.try_acquire("key1", tokens=10.0) is True 

110 assert tb.try_acquire("key1") is False 

111 

112 

113class TestTokenBucketAcquireOrWait: 

114 def test_immediate_success(self): 

115 tb = TokenBucket(rate=10.0, burst=5.0) 

116 assert tb.acquire_or_wait("key1", timeout=1.0) is True 

117 

118 def test_blocks_and_waits_for_refill(self): 

119 tb = TokenBucket(rate=200.0, burst=2.0) # 200 tokens/sec 

120 for _ in range(2): 

121 tb.try_acquire("key1") 

122 # Should block briefly then succeed 

123 assert tb.acquire_or_wait("key1", timeout=1.0) is True 

124 

125 def test_timeout_expires(self): 

126 tb = TokenBucket(rate=0.1, burst=1.0) # very slow refill 

127 tb.try_acquire("key1") 

128 assert tb.acquire_or_wait("key1", timeout=0.01) is False 

129 

130 def test_no_timeout_blocks_indefinitely(self): 

131 tb = TokenBucket(rate=500.0, burst=1.0) 

132 tb.try_acquire("key1") 

133 # Should succeed within a very short time without explicit timeout 

134 assert tb.acquire_or_wait("key1", timeout=None) is True 

135 

136 def test_multi_token_acquire_or_wait(self): 

137 tb = TokenBucket(rate=500.0, burst=5.0) 

138 assert tb.acquire_or_wait("key1", timeout=1.0, tokens=5.0) is True 

139 assert tb.try_acquire("key1") is False 

140 

141 

142class TestTokenBucketReset: 

143 def test_reset_removes_key(self): 

144 tb = TokenBucket(rate=10.0, burst=3.0) 

145 for _ in range(3): 

146 tb.try_acquire("key1") 

147 assert tb.try_acquire("key1") is False 

148 tb.reset("key1") 

149 assert tb.try_acquire("key1") is True 

150 

151 def test_reset_nonexistent_key_no_error(self): 

152 tb = TokenBucket(rate=10.0) 

153 tb.reset("no-such-key") 

154 

155 def test_reset_all_clears_everything(self): 

156 tb = TokenBucket(rate=10.0, burst=2.0) 

157 tb.try_acquire("key1") 

158 tb.try_acquire("key2") 

159 tb.reset_all() 

160 assert len(tb._buckets) == 0 

161 assert tb.try_acquire("key1") is True 

162 

163 

164class TestTokenBucketStats: 

165 def test_initial_stats(self): 

166 tb = TokenBucket(rate=10.0, burst=20.0) 

167 s = tb.stats() 

168 assert s["rate"] == 10.0 

169 assert s["burst"] == 20.0 

170 assert s["active_keys"] == 0 

171 assert s["total_acquired"] == 0 

172 assert s["total_rejected"] == 0 

173 

174 def test_stats_after_acquire(self): 

175 tb = TokenBucket(rate=10.0, burst=5.0) 

176 for _ in range(3): 

177 tb.try_acquire("key1") 

178 s = tb.stats() 

179 assert s["total_acquired"] == 3 

180 assert s["active_keys"] == 1 

181 

182 def test_stats_after_reject(self): 

183 tb = TokenBucket(rate=10.0, burst=1.0) 

184 tb.try_acquire("key1") 

185 tb.try_acquire("key1") # rejected — same key, bucket exhausted 

186 s = tb.stats() 

187 assert s["total_acquired"] == 1 

188 assert s["total_rejected"] == 1 

189 

190 

191# ============================================================================ 

192# SlidingWindow 

193# ============================================================================ 

194 

195 

196class TestSlidingWindowInit: 

197 def test_default_window(self): 

198 sw = SlidingWindow(limit=100) 

199 assert sw._limit == 100 

200 assert sw._window == 60.0 

201 

202 def test_custom_window(self): 

203 sw = SlidingWindow(limit=50, window=30.0) 

204 assert sw._limit == 50 

205 assert sw._window == 30.0 

206 

207 def test_limit_zero_raises(self): 

208 with pytest.raises(ValueError, match="positive"): 

209 SlidingWindow(limit=0) 

210 

211 def test_limit_negative_raises(self): 

212 with pytest.raises(ValueError, match="positive"): 

213 SlidingWindow(limit=-5) 

214 

215 def test_limit_property(self): 

216 sw = SlidingWindow(limit=42) 

217 assert sw.limit == 42 

218 

219 

220class TestSlidingWindowTryAcquire: 

221 def test_within_limit(self): 

222 sw = SlidingWindow(limit=5, window=60.0) 

223 for _ in range(5): 

224 assert sw.try_acquire("key1") is True 

225 

226 def test_exceeds_limit(self): 

227 sw = SlidingWindow(limit=3, window=60.0) 

228 for _ in range(3): 

229 assert sw.try_acquire("key1") is True 

230 assert sw.try_acquire("key1") is False 

231 

232 def test_different_keys_independent(self): 

233 sw = SlidingWindow(limit=2, window=60.0) 

234 assert sw.try_acquire("key1") is True 

235 assert sw.try_acquire("key1") is True 

236 assert sw.try_acquire("key1") is False 

237 assert sw.try_acquire("key2") is True 

238 

239 def test_eviction_of_old_entries(self): 

240 sw = SlidingWindow(limit=2, window=0.05) # 50ms window 

241 assert sw.try_acquire("key1") is True 

242 assert sw.try_acquire("key1") is True 

243 assert sw.try_acquire("key1") is False 

244 time.sleep(0.06) 

245 assert sw.try_acquire("key1") is True 

246 

247 

248class TestSlidingWindowAcquireOrWait: 

249 def test_immediate_success(self): 

250 sw = SlidingWindow(limit=5, window=60.0) 

251 assert sw.acquire_or_wait("key1", timeout=1.0) is True 

252 

253 def test_blocks_and_waits(self): 

254 sw = SlidingWindow(limit=2, window=0.05) 

255 for _ in range(2): 

256 sw.try_acquire("key1") 

257 assert sw.acquire_or_wait("key1", timeout=1.0) is True 

258 

259 def test_timeout_expires(self): 

260 sw = SlidingWindow(limit=1, window=60.0) 

261 sw.try_acquire("key1") 

262 assert sw.acquire_or_wait("key1", timeout=0.01) is False 

263 

264 def test_no_timeout_succeeds(self): 

265 sw = SlidingWindow(limit=1, window=0.02) 

266 sw.try_acquire("key1") 

267 assert sw.acquire_or_wait("key1", timeout=None) is True 

268 

269 

270class TestSlidingWindowReset: 

271 def test_reset_removes_key(self): 

272 sw = SlidingWindow(limit=1, window=60.0) 

273 sw.try_acquire("key1") 

274 assert sw.try_acquire("key1") is False 

275 sw.reset("key1") 

276 assert sw.try_acquire("key1") is True 

277 

278 def test_reset_nonexistent_no_error(self): 

279 sw = SlidingWindow(limit=10) 

280 sw.reset("ghost") 

281 

282 def test_reset_all(self): 

283 sw = SlidingWindow(limit=2, window=60.0) 

284 sw.try_acquire("key1") 

285 sw.try_acquire("key2") 

286 sw.reset_all() 

287 assert len(sw._windows) == 0 

288 assert sw.try_acquire("key1") is True 

289 

290 

291class TestSlidingWindowStats: 

292 def test_initial_stats(self): 

293 sw = SlidingWindow(limit=100, window=60.0) 

294 s = sw.stats() 

295 assert s["limit"] == 100 

296 assert s["window"] == 60.0 

297 assert s["active_keys"] == 0 

298 assert s["total_acquired"] == 0 

299 assert s["total_rejected"] == 0 

300 

301 def test_stats_after_use(self): 

302 sw = SlidingWindow(limit=5, window=60.0) 

303 for _ in range(3): 

304 sw.try_acquire("key1") 

305 sw.try_acquire("key2") 

306 s = sw.stats() 

307 assert s["total_acquired"] == 4 

308 assert s["active_keys"] >= 1 

309 

310 def test_stats_after_reject(self): 

311 sw = SlidingWindow(limit=2, window=60.0) 

312 for _ in range(2): 

313 sw.try_acquire("key1") 

314 sw.try_acquire("key1") # rejected 

315 sw.try_acquire("key1") # rejected 

316 s = sw.stats() 

317 assert s["total_acquired"] == 2 

318 assert s["total_rejected"] == 2