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# ============================================================================ 

15# RateLimitExceeded 

16# ============================================================================ 

17 

18 

19class TestRateLimitExceeded: 

20 def test_constructor_and_attributes(self): 

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

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

23 assert exc.limit == 10.0 

24 assert exc.window == 60.0 

25 

26 def test_string_message(self): 

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

28 msg = str(exc) 

29 assert "user:42" in msg 

30 assert "5.0" in msg 

31 assert "1.0" in msg 

32 

33 

34# ============================================================================ 

35# TokenBucket 

36# ============================================================================ 

37 

38 

39class TestTokenBucketInit: 

40 def test_default_burst_equals_rate(self): 

41 tb = TokenBucket(rate=10.0) 

42 assert tb._rate == 10.0 

43 assert tb._burst == 10.0 

44 

45 def test_custom_burst(self): 

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

47 assert tb._rate == 5.0 

48 assert tb._burst == 20.0 

49 

50 def test_rate_zero_raises(self): 

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

52 TokenBucket(rate=0) 

53 

54 def test_rate_negative_raises(self): 

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

56 TokenBucket(rate=-1.0) 

57 

58 def test_rate_property(self): 

59 tb = TokenBucket(rate=7.5) 

60 assert tb.rate == 7.5 

61 

62 

63class TestTokenBucketTryAcquire: 

64 def test_first_acquire_succeeds(self): 

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

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

67 

68 def test_acquire_within_burst(self): 

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

70 for _ in range(5): 

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

72 # 6th should fail 

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

74 

75 def test_acquire_multi_tokens(self): 

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

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

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

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

80 

81 def test_different_keys_independent(self): 

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

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

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

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

86 # key2 should be unaffected 

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

88 

89 def test_refill_after_wait(self): 

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

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

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

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

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

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

96 

97 def test_refill_capped_at_burst(self): 

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

99 for _ in range(3): 

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

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

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

103 # should have exactly 3 tokens, not 100 

104 for _ in range(3): 

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

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

107 

108 def test_acquire_exact_tokens(self): 

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

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

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

112 

113 

114class TestTokenBucketAcquireOrWait: 

115 def test_immediate_success(self): 

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

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

118 

119 def test_blocks_and_waits_for_refill(self): 

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

121 for _ in range(2): 

122 tb.try_acquire("key1") 

123 # Should block briefly then succeed 

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

125 

126 def test_timeout_expires(self): 

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

128 tb.try_acquire("key1") 

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

130 

131 def test_no_timeout_blocks_indefinitely(self): 

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

133 tb.try_acquire("key1") 

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

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

136 

137 def test_multi_token_acquire_or_wait(self): 

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

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

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

141 

142 

143class TestTokenBucketReset: 

144 def test_reset_removes_key(self): 

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

146 for _ in range(3): 

147 tb.try_acquire("key1") 

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

149 tb.reset("key1") 

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

151 

152 def test_reset_nonexistent_key_no_error(self): 

153 tb = TokenBucket(rate=10.0) 

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

155 

156 def test_reset_all_clears_everything(self): 

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

158 tb.try_acquire("key1") 

159 tb.try_acquire("key2") 

160 tb.reset_all() 

161 assert len(tb._buckets) == 0 

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

163 

164 

165class TestTokenBucketStats: 

166 def test_initial_stats(self): 

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

168 s = tb.stats() 

169 assert s["rate"] == 10.0 

170 assert s["burst"] == 20.0 

171 assert s["active_keys"] == 0 

172 assert s["total_acquired"] == 0 

173 assert s["total_rejected"] == 0 

174 

175 def test_stats_after_acquire(self): 

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

177 for _ in range(3): 

178 tb.try_acquire("key1") 

179 s = tb.stats() 

180 assert s["total_acquired"] == 3 

181 assert s["active_keys"] == 1 

182 

183 def test_stats_after_reject(self): 

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

185 tb.try_acquire("key1") 

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

187 s = tb.stats() 

188 assert s["total_acquired"] == 1 

189 assert s["total_rejected"] == 1 

190 

191 

192# ============================================================================ 

193# SlidingWindow 

194# ============================================================================ 

195 

196 

197class TestSlidingWindowInit: 

198 def test_default_window(self): 

199 sw = SlidingWindow(limit=100) 

200 assert sw._limit == 100 

201 assert sw._window == 60.0 

202 

203 def test_custom_window(self): 

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

205 assert sw._limit == 50 

206 assert sw._window == 30.0 

207 

208 def test_limit_zero_raises(self): 

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

210 SlidingWindow(limit=0) 

211 

212 def test_limit_negative_raises(self): 

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

214 SlidingWindow(limit=-5) 

215 

216 def test_limit_property(self): 

217 sw = SlidingWindow(limit=42) 

218 assert sw.limit == 42 

219 

220 

221class TestSlidingWindowTryAcquire: 

222 def test_within_limit(self): 

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

224 for _ in range(5): 

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

226 

227 def test_exceeds_limit(self): 

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

229 for _ in range(3): 

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

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

232 

233 def test_different_keys_independent(self): 

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

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

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

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

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

239 

240 def test_eviction_of_old_entries(self): 

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

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

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

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

245 time.sleep(0.06) 

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

247 

248 

249class TestSlidingWindowAcquireOrWait: 

250 def test_immediate_success(self): 

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

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

253 

254 def test_blocks_and_waits(self): 

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

256 for _ in range(2): 

257 sw.try_acquire("key1") 

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

259 

260 def test_timeout_expires(self): 

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

262 sw.try_acquire("key1") 

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

264 

265 def test_no_timeout_succeeds(self): 

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

267 sw.try_acquire("key1") 

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

269 

270 

271class TestSlidingWindowReset: 

272 def test_reset_removes_key(self): 

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

274 sw.try_acquire("key1") 

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

276 sw.reset("key1") 

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

278 

279 def test_reset_nonexistent_no_error(self): 

280 sw = SlidingWindow(limit=10) 

281 sw.reset("ghost") 

282 

283 def test_reset_all(self): 

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

285 sw.try_acquire("key1") 

286 sw.try_acquire("key2") 

287 sw.reset_all() 

288 assert len(sw._windows) == 0 

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

290 

291 

292class TestSlidingWindowStats: 

293 def test_initial_stats(self): 

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

295 s = sw.stats() 

296 assert s["limit"] == 100 

297 assert s["window"] == 60.0 

298 assert s["active_keys"] == 0 

299 assert s["total_acquired"] == 0 

300 assert s["total_rejected"] == 0 

301 

302 def test_stats_after_use(self): 

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

304 for _ in range(3): 

305 sw.try_acquire("key1") 

306 sw.try_acquire("key2") 

307 s = sw.stats() 

308 assert s["total_acquired"] == 4 

309 assert s["active_keys"] >= 1 

310 

311 def test_stats_after_reject(self): 

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

313 for _ in range(2): 

314 sw.try_acquire("key1") 

315 sw.try_acquire("key1") # rejected 

316 sw.try_acquire("key1") # rejected 

317 s = sw.stats() 

318 assert s["total_acquired"] == 2 

319 assert s["total_rejected"] == 2