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

1"""Tests for agentos.core.rate_limiter — TokenBucket, SlidingWindow, ConcurrentLimiter, etc.""" 

2 

3import asyncio 

4import time 

5import pytest 

6from agentos.core.rate_limiter import ( 

7 CompositeLimiter, 

8 ConcurrentLimiter, 

9 EndpointRateLimit, 

10 RateLimitError, 

11 RateLimiter, 

12 RateLimitRegistry, 

13 SlidingWindow, 

14 TokenBucket, 

15) 

16 

17 

18# ============================================================================ 

19# TokenBucket 

20# ============================================================================ 

21 

22class TestTokenBucket: 

23 def test_init_defaults(self): 

24 tb = TokenBucket(rate=10.0, capacity=100) 

25 assert tb.rate == 10.0 

26 assert tb.capacity == 100 

27 

28 def test_invalid_rate(self): 

29 with pytest.raises(ValueError, match="rate"): 

30 TokenBucket(rate=0, capacity=10) 

31 with pytest.raises(ValueError, match="rate"): 

32 TokenBucket(rate=-1, capacity=10) 

33 

34 def test_invalid_capacity(self): 

35 with pytest.raises(ValueError, match="capacity"): 

36 TokenBucket(rate=10, capacity=0) 

37 with pytest.raises(ValueError, match="capacity"): 

38 TokenBucket(rate=10, capacity=-5) 

39 

40 @pytest.mark.asyncio 

41 async def test_acquire_initial_tokens(self): 

42 tb = TokenBucket(rate=1.0, capacity=5) 

43 assert await tb.acquire() is True 

44 assert await tb.acquire() is True 

45 assert await tb.acquire() is True 

46 

47 @pytest.mark.asyncio 

48 async def test_acquire_exhausted(self): 

49 tb = TokenBucket(rate=0.001, capacity=2) 

50 await tb.acquire() 

51 await tb.acquire() 

52 assert await tb.acquire() is False 

53 

54 @pytest.mark.asyncio 

55 async def test_refill_over_time(self): 

56 tb = TokenBucket(rate=100.0, capacity=10) 

57 # Use all tokens 

58 for _ in range(10): 

59 assert await tb.acquire() 

60 assert await tb.acquire() is False 

61 

62 # Wait for refill 

63 await asyncio.sleep(0.05) 

64 assert await tb.acquire() is True 

65 

66 @pytest.mark.asyncio 

67 async def test_wait_and_acquire(self): 

68 tb = TokenBucket(rate=200.0, capacity=2) 

69 for _ in range(2): 

70 await tb.acquire() 

71 

72 # Should wait briefly and succeed 

73 result = await tb.wait_and_acquire(timeout=2.0) 

74 assert result is True 

75 

76 @pytest.mark.asyncio 

77 async def test_wait_and_acquire_timeout(self): 

78 tb = TokenBucket(rate=0.1, capacity=1) 

79 await tb.acquire() 

80 

81 result = await tb.wait_and_acquire(timeout=0.01) 

82 assert result is False 

83 

84 @pytest.mark.asyncio 

85 async def test_available_tokens(self): 

86 tb = TokenBucket(rate=1.0, capacity=100) 

87 assert tb.available_tokens == 100.0 

88 await tb.acquire(tokens=30) 

89 assert tb.available_tokens == 70.0 

90 

91 def test_fill_level(self): 

92 tb = TokenBucket(rate=1.0, capacity=100) 

93 assert tb.fill_level == 1.0 

94 

95 

96# ============================================================================ 

97# SlidingWindow 

98# ============================================================================ 

99 

100class TestSlidingWindow: 

101 def test_init(self): 

102 sw = SlidingWindow(max_requests=5, window_seconds=10.0) 

103 assert sw.max_requests == 5 

104 assert sw.window_seconds == 10.0 

105 

106 def test_invalid_max_requests(self): 

107 with pytest.raises(ValueError): 

108 SlidingWindow(max_requests=0, window_seconds=1) 

109 

110 def test_invalid_window(self): 

111 with pytest.raises(ValueError): 

112 SlidingWindow(max_requests=5, window_seconds=0) 

113 

114 @pytest.mark.asyncio 

115 async def test_acquire_within_limit(self): 

116 sw = SlidingWindow(max_requests=3, window_seconds=60) 

117 assert await sw.acquire() is True 

118 assert await sw.acquire() is True 

119 assert await sw.acquire() is True 

120 

121 @pytest.mark.asyncio 

122 async def test_acquire_exceeds_limit(self): 

123 sw = SlidingWindow(max_requests=2, window_seconds=60) 

124 await sw.acquire() 

125 await sw.acquire() 

126 assert await sw.acquire() is False 

127 

128 @pytest.mark.asyncio 

129 async def test_current_count(self): 

130 sw = SlidingWindow(max_requests=10, window_seconds=60) 

131 await sw.acquire() 

132 await sw.acquire() 

133 assert sw.current_count == 2 

134 

135 def test_remaining(self): 

136 sw = SlidingWindow(max_requests=10, window_seconds=60) 

137 assert sw.remaining == 10 

138 

139 

140# ============================================================================ 

141# ConcurrentLimiter 

142# ============================================================================ 

143 

144class TestConcurrentLimiter: 

145 def test_init(self): 

146 cl = ConcurrentLimiter(max_concurrent=5) 

147 assert cl.available == 5 

148 

149 def test_invalid_max(self): 

150 with pytest.raises(ValueError): 

151 ConcurrentLimiter(max_concurrent=0) 

152 

153 @pytest.mark.asyncio 

154 async def test_acquire_release(self): 

155 cl = ConcurrentLimiter(max_concurrent=2) 

156 assert await cl.acquire() is True 

157 assert await cl.acquire() is True 

158 assert cl.available == 0 

159 cl.release() 

160 assert cl.available == 1 

161 

162 @pytest.mark.asyncio 

163 async def test_acquire_blocking(self): 

164 cl = ConcurrentLimiter(max_concurrent=1) 

165 await cl.acquire() 

166 

167 # Acquire in task that will block 

168 async def blocked(): 

169 return await cl.acquire() 

170 

171 task = asyncio.create_task(blocked()) 

172 await asyncio.sleep(0.1) 

173 assert not task.done() 

174 

175 cl.release() 

176 result = await asyncio.wait_for(task, timeout=2.0) 

177 assert result is True 

178 

179 

180# ============================================================================ 

181# CompositeLimiter 

182# ============================================================================ 

183 

184class TestCompositeLimiter: 

185 @pytest.mark.asyncio 

186 async def test_all_pass(self): 

187 tb = TokenBucket(rate=100, capacity=10) 

188 sw = SlidingWindow(max_requests=10, window_seconds=60) 

189 cl = ConcurrentLimiter(max_concurrent=5) 

190 composite = CompositeLimiter([tb, sw, cl]) 

191 assert await composite.acquire() is True 

192 

193 @pytest.mark.asyncio 

194 async def test_one_fails(self): 

195 tb = TokenBucket(rate=0.001, capacity=1) 

196 sw = SlidingWindow(max_requests=10, window_seconds=60) 

197 await tb.acquire() # exhaust token bucket 

198 composite = CompositeLimiter([tb, sw]) 

199 assert await composite.acquire() is False 

200 

201 

202# ============================================================================ 

203# RateLimitError 

204# ============================================================================ 

205 

206class TestRateLimitError: 

207 def test_message(self): 

208 err = RateLimitError("Limit exceeded: api") 

209 assert "api" in str(err) 

210 

211 

212# ============================================================================ 

213# RateLimiter — context manager 

214# ============================================================================ 

215 

216class TestRateLimiterCM: 

217 @pytest.mark.asyncio 

218 async def test_async_context_manager_success(self): 

219 rl = RateLimiter(name="test", strategy=TokenBucket(rate=10, capacity=10)) 

220 async with rl: 

221 pass 

222 

223 @pytest.mark.asyncio 

224 async def test_async_context_manager_rejected(self): 

225 tb = TokenBucket(rate=0.001, capacity=1) 

226 await tb.acquire() 

227 rl = RateLimiter(name="test", strategy=tb) 

228 with pytest.raises(RateLimitError): 

229 async with rl: 

230 pass 

231 

232 @pytest.mark.asyncio 

233 async def test_concurrent_releases_on_exit(self): 

234 cl = ConcurrentLimiter(max_concurrent=2) 

235 rl = RateLimiter(name="test", strategy=cl) 

236 async with rl: 

237 assert cl.available == 1 

238 assert cl.available == 2 # released on exit 

239 

240 

241# ============================================================================ 

242# RateLimiter — factory methods 

243# ============================================================================ 

244 

245class TestRateLimiterFactories: 

246 def test_token_bucket_factory(self): 

247 rl = RateLimiter.token_bucket("api", rate=5.0, capacity=20) 

248 assert rl.name == "api" 

249 assert isinstance(rl.strategy, TokenBucket) 

250 assert rl.strategy.rate == 5.0 

251 assert rl.strategy.capacity == 20 

252 

253 def test_sliding_window_factory(self): 

254 rl = RateLimiter.sliding_window("api", max_requests=100, window_seconds=60) 

255 assert rl.name == "api" 

256 assert isinstance(rl.strategy, SlidingWindow) 

257 assert rl.strategy.max_requests == 100 

258 

259 def test_concurrent_factory(self): 

260 rl = RateLimiter.concurrent("api", max_concurrent=10) 

261 assert rl.name == "api" 

262 assert isinstance(rl.strategy, ConcurrentLimiter) 

263 

264 

265# ============================================================================ 

266# EndpointRateLimit 

267# ============================================================================ 

268 

269class TestEndpointRateLimit: 

270 def test_defaults(self): 

271 erl = EndpointRateLimit(endpoint="/api/v1") 

272 assert erl.endpoint == "/api/v1" 

273 assert erl.requests_per_second is None 

274 assert erl.concurrent is None 

275 assert erl.burst == 1 

276 

277 def test_full_spec(self): 

278 erl = EndpointRateLimit( 

279 endpoint="/api/v1", 

280 requests_per_second=10.0, 

281 requests_per_minute=600, 

282 concurrent=5, 

283 burst=20, 

284 ) 

285 assert erl.requests_per_second == 10.0 

286 assert erl.requests_per_minute == 600 

287 assert erl.concurrent == 5 

288 assert erl.burst == 20 

289 

290 

291# ============================================================================ 

292# RateLimitRegistry 

293# ============================================================================ 

294 

295class TestRateLimitRegistry: 

296 @pytest.mark.asyncio 

297 async def test_configure_single_limiter(self): 

298 reg = RateLimitRegistry() 

299 spec = EndpointRateLimit(endpoint="/api", requests_per_second=10.0, burst=5) 

300 limiter = await reg.configure(spec) 

301 assert limiter.name == "/api" 

302 assert isinstance(limiter.strategy, TokenBucket) 

303 

304 @pytest.mark.asyncio 

305 async def test_configure_composite(self): 

306 reg = RateLimitRegistry() 

307 spec = EndpointRateLimit( 

308 endpoint="/api", 

309 requests_per_second=10.0, 

310 requests_per_minute=600, 

311 concurrent=5, 

312 ) 

313 limiter = await reg.configure(spec) 

314 assert limiter.name == "/api" 

315 assert isinstance(limiter.strategy, CompositeLimiter) 

316 

317 @pytest.mark.asyncio 

318 async def test_get_existing(self): 

319 reg = RateLimitRegistry() 

320 await reg.configure(EndpointRateLimit(endpoint="/api", requests_per_second=1.0)) 

321 limiter = await reg.get("/api") 

322 assert limiter is not None 

323 assert limiter.name == "/api" 

324 

325 @pytest.mark.asyncio 

326 async def test_get_missing(self): 

327 reg = RateLimitRegistry() 

328 assert await reg.get("/nonexistent") is None 

329 

330 @pytest.mark.asyncio 

331 async def test_acquire_existing(self): 

332 reg = RateLimitRegistry() 

333 await reg.configure(EndpointRateLimit(endpoint="/api", requests_per_second=100.0, burst=100)) 

334 result = await reg.acquire("/api") 

335 assert result is True 

336 

337 @pytest.mark.asyncio 

338 async def test_acquire_missing(self): 

339 reg = RateLimitRegistry() 

340 result = await reg.acquire("/nonexistent") 

341 assert result is True 

342 

343 @pytest.mark.asyncio 

344 async def test_configure_with_burst_one(self): 

345 reg = RateLimitRegistry() 

346 spec = EndpointRateLimit(endpoint="/api", requests_per_second=5.0) 

347 limiter = await reg.configure(spec) 

348 assert isinstance(limiter.strategy, TokenBucket) 

349 assert limiter.strategy.capacity == 1 # burst=1 default