Coverage for agentos/tests/test_circuit_breaker.py: 99%

1"""Tests for agentos.core.circuit_breaker — CircuitBreaker, CircuitRegistry, decorator.""" 

2 

3import asyncio 

4import time 

5import pytest 

6from agentos.core.circuit_breaker import ( 

7 CircuitBreaker, 

8 CircuitConfig, 

9 CircuitOpenError, 

10 CircuitRegistry, 

11 CircuitState, 

12 CircuitStats, 

13 circuit_breaker, 

14 default_registry, 

15) 

16 

17 

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

19# CircuitState 

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

21 

22class TestCircuitState: 

23 def test_enum_values(self): 

24 assert CircuitState.CLOSED == "closed" 

25 assert CircuitState.OPEN == "open" 

26 assert CircuitState.HALF_OPEN == "half_open" 

27 

28 

29# ============================================================================ 

30# CircuitConfig 

31# ============================================================================ 

32 

33class TestCircuitConfig: 

34 def test_defaults(self): 

35 cfg = CircuitConfig() 

36 assert cfg.failure_threshold == 5 

37 assert cfg.success_threshold == 2 

38 assert cfg.timeout_seconds == 60.0 

39 assert cfg.half_open_max_requests == 1 

40 assert cfg.excluded_exceptions == () 

41 

42 def test_custom(self): 

43 cfg = CircuitConfig( 

44 failure_threshold=3, 

45 success_threshold=1, 

46 timeout_seconds=10.0, 

47 ) 

48 assert cfg.failure_threshold == 3 

49 assert cfg.timeout_seconds == 10.0 

50 

51 

52# ============================================================================ 

53# CircuitStats 

54# ============================================================================ 

55 

56class TestCircuitStats: 

57 def test_defaults(self): 

58 s = CircuitStats() 

59 assert s.state == CircuitState.CLOSED 

60 assert s.failure_count == 0 

61 assert s.total_failures == 0 

62 

63 def test_reset(self): 

64 s = CircuitStats() 

65 s.failure_count = 5 

66 s.success_count = 3 

67 s.half_open_requests = 2 

68 s.reset() 

69 assert s.failure_count == 0 

70 assert s.success_count == 0 

71 assert s.half_open_requests == 0 

72 

73 

74# ============================================================================ 

75# CircuitBreaker — Basic 

76# ============================================================================ 

77 

78class TestCircuitBreakerBasic: 

79 def test_default_state(self): 

80 cb = CircuitBreaker(name="test") 

81 assert cb.state == CircuitState.CLOSED 

82 assert cb.name == "test" 

83 

84 def test_custom_config(self): 

85 cfg = CircuitConfig(failure_threshold=2) 

86 cb = CircuitBreaker(name="test", config=cfg) 

87 assert cb.config.failure_threshold == 2 

88 

89 

90# ============================================================================ 

91# CircuitBreaker — Acquire / Release 

92# ============================================================================ 

93 

94class TestCircuitBreakerAcquire: 

95 @pytest.mark.asyncio 

96 async def test_closed_always_acquires(self): 

97 cb = CircuitBreaker(name="test") 

98 assert await cb.acquire() is True 

99 assert await cb.acquire() is True 

100 

101 @pytest.mark.asyncio 

102 async def test_release_half_open(self): 

103 cb = CircuitBreaker(name="test") 

104 # Force half-open state for testing release 

105 cb.stats.state = CircuitState.HALF_OPEN 

106 cb.stats.half_open_requests = 1 

107 await cb.release() 

108 assert cb.stats.half_open_requests == 0 

109 

110 @pytest.mark.asyncio 

111 async def test_release_floor_zero(self): 

112 cb = CircuitBreaker(name="test") 

113 cb.stats.state = CircuitState.HALF_OPEN 

114 cb.stats.half_open_requests = 0 

115 await cb.release() 

116 assert cb.stats.half_open_requests == 0 

117 

118 

119# ============================================================================ 

120# CircuitBreaker — State transitions 

121# ============================================================================ 

122 

123class TestCircuitBreakerTransitions: 

124 @pytest.mark.asyncio 

125 async def test_success_keeps_closed(self): 

126 cb = CircuitBreaker(name="test") 

127 

128 async def ok(): return "success" 

129 result = await cb.call(ok) 

130 assert result == "success" 

131 assert cb.state == CircuitState.CLOSED 

132 

133 @pytest.mark.asyncio 

134 async def test_failure_below_threshold(self): 

135 cb = CircuitBreaker(name="test", config=CircuitConfig(failure_threshold=3)) 

136 

137 async def fail(): 

138 raise ValueError("fail") 

139 

140 for _ in range(2): 

141 with pytest.raises(ValueError): 

142 await cb.call(fail) 

143 

144 assert cb.state == CircuitState.CLOSED 

145 

146 @pytest.mark.asyncio 

147 async def test_failure_trips_open(self): 

148 cb = CircuitBreaker(name="test", config=CircuitConfig(failure_threshold=2)) 

149 

150 async def fail(): 

151 raise ValueError("boom") 

152 

153 for _ in range(2): 

154 with pytest.raises(ValueError): 

155 await cb.call(fail) 

156 

157 assert cb.state == CircuitState.OPEN 

158 

159 @pytest.mark.asyncio 

160 async def test_open_rejects_requests(self): 

161 cb = CircuitBreaker(name="test", config=CircuitConfig(failure_threshold=1, timeout_seconds=60)) 

162 

163 async def fail(): 

164 raise ValueError("boom") 

165 

166 with pytest.raises(ValueError): 

167 await cb.call(fail) 

168 

169 assert cb.state == CircuitState.OPEN 

170 # Subsequent request should raise CircuitOpenError 

171 with pytest.raises(CircuitOpenError): 

172 await cb.call(lambda: "ok") 

173 

174 @pytest.mark.asyncio 

175 async def test_half_open_probing(self): 

176 cb = CircuitBreaker( 

177 name="test", 

178 config=CircuitConfig(failure_threshold=1, timeout_seconds=0.01, success_threshold=1), 

179 ) 

180 

181 async def fail(): 

182 raise ValueError("boom") 

183 

184 # Trip open 

185 with pytest.raises(ValueError): 

186 await cb.call(fail) 

187 assert cb.state == CircuitState.OPEN 

188 

189 # Wait for timeout to expire 

190 await asyncio.sleep(0.05) 

191 

192 # Now should probe (half-open) 

193 async def ok(): 

194 return "recovered" 

195 

196 result = await cb.call(ok) 

197 assert result == "recovered" 

198 assert cb.state == CircuitState.CLOSED 

199 

200 @pytest.mark.asyncio 

201 async def test_half_open_failure_reopens(self): 

202 cb = CircuitBreaker( 

203 name="test", 

204 config=CircuitConfig( 

205 failure_threshold=1, 

206 timeout_seconds=0.01, 

207 success_threshold=1, 

208 ), 

209 ) 

210 

211 async def fail(): 

212 raise ValueError("boom") 

213 

214 # Trip open 

215 with pytest.raises(ValueError): 

216 await cb.call(fail) 

217 assert cb.state == CircuitState.OPEN 

218 

219 await asyncio.sleep(0.05) 

220 

221 # Probe fails → back to OPEN 

222 with pytest.raises(ValueError): 

223 await cb.call(fail) 

224 assert cb.state == CircuitState.OPEN 

225 

226 @pytest.mark.asyncio 

227 async def test_success_resets_count_in_closed(self): 

228 cb = CircuitBreaker(name="test", config=CircuitConfig(failure_threshold=3)) 

229 

230 async def fail(): 

231 raise ValueError("boom") 

232 

233 async def ok(): 

234 return "ok" 

235 

236 with pytest.raises(ValueError): 

237 await cb.call(fail) 

238 await cb.call(ok) 

239 # One failure then success should reset failure count 

240 assert cb.stats.failure_count == 0 

241 

242 

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

244# CircuitBreaker — Excluded exceptions 

245# ============================================================================ 

246 

247class TestCircuitBreakerExcluded: 

248 @pytest.mark.asyncio 

249 async def test_excluded_exception_does_not_count(self): 

250 cfg = CircuitConfig(failure_threshold=1, excluded_exceptions=(ValueError,)) 

251 cb = CircuitBreaker(name="test", config=cfg) 

252 

253 async def fail(): 

254 raise ValueError("ignored") 

255 

256 # This should NOT count toward failure threshold 

257 with pytest.raises(ValueError): 

258 await cb.call(fail) 

259 assert cb.state == CircuitState.CLOSED 

260 

261 

262# ============================================================================ 

263# CircuitBreaker — Fallback 

264# ============================================================================ 

265 

266class TestCircuitBreakerFallback: 

267 @pytest.mark.asyncio 

268 async def test_fallback_used_when_open(self): 

269 cb = CircuitBreaker( 

270 name="test", 

271 config=CircuitConfig(failure_threshold=1, timeout_seconds=60), 

272 ) 

273 

274 async def fail(): 

275 raise ValueError("boom") 

276 

277 with pytest.raises(ValueError): 

278 await cb.call(fail) 

279 

280 assert cb.state == CircuitState.OPEN 

281 

282 # Use fallback 

283 async def fallback_fn(): 

284 return "fallback_value" 

285 

286 result = await cb.call(lambda: "ok", fallback=fallback_fn) 

287 assert result == "fallback_value" 

288 

289 

290# ============================================================================ 

291# CircuitBreaker — config reset behavior 

292# ============================================================================ 

293 

294class TestCircuitBreakerConfig: 

295 @pytest.mark.asyncio 

296 async def test_failure_threshold_one(self): 

297 cb = CircuitBreaker(name="test", config=CircuitConfig(failure_threshold=1)) 

298 

299 async def fail(): 

300 raise ValueError("x") 

301 

302 with pytest.raises(ValueError): 

303 await cb.call(fail) 

304 assert cb.state == CircuitState.OPEN 

305 

306 

307# ============================================================================ 

308# CircuitRegistry 

309# ============================================================================ 

310 

311class TestCircuitRegistry: 

312 @pytest.mark.asyncio 

313 async def test_get_or_create(self): 

314 reg = CircuitRegistry() 

315 cb1 = await reg.get_or_create("api") 

316 cb2 = await reg.get_or_create("api") 

317 assert cb1 is cb2 

318 

319 @pytest.mark.asyncio 

320 async def test_get_or_create_different_names(self): 

321 reg = CircuitRegistry() 

322 cb1 = await reg.get_or_create("a") 

323 cb2 = await reg.get_or_create("b") 

324 assert cb1 is not cb2 

325 

326 @pytest.mark.asyncio 

327 async def test_get_or_create_with_config(self): 

328 reg = CircuitRegistry() 

329 cfg = CircuitConfig(failure_threshold=10) 

330 cb = await reg.get_or_create("api", config=cfg) 

331 assert cb.config.failure_threshold == 10 

332 

333 @pytest.mark.asyncio 

334 async def test_get_all_stats(self): 

335 reg = CircuitRegistry() 

336 await reg.get_or_create("a") 

337 await reg.get_or_create("b") 

338 stats = reg.get_all_stats() 

339 assert "a" in stats 

340 assert "b" in stats 

341 

342 @pytest.mark.asyncio 

343 async def test_reset_all(self): 

344 reg = CircuitRegistry() 

345 cb = await reg.get_or_create("test") 

346 

347 async def fail(): 

348 raise ValueError("x") 

349 cb_cfg = CircuitConfig(failure_threshold=1, timeout_seconds=60) 

350 cb.config = cb_cfg 

351 with pytest.raises(ValueError): 

352 await cb.call(fail) 

353 assert cb.state == CircuitState.OPEN 

354 

355 await reg.reset_all() 

356 # After reset, failure_count is 0 and state is CLOSED 

357 assert cb.stats.failure_count == 0 

358 

359 @pytest.mark.asyncio 

360 async def test_force_open(self): 

361 reg = CircuitRegistry() 

362 cb = await reg.get_or_create("test") 

363 await reg.force_open("test") 

364 assert cb.state == CircuitState.OPEN 

365 

366 @pytest.mark.asyncio 

367 async def test_force_closed(self): 

368 reg = CircuitRegistry() 

369 cb = await reg.get_or_create("test") 

370 cb.stats.state = CircuitState.OPEN 

371 await reg.force_closed("test") 

372 assert cb.state == CircuitState.CLOSED 

373 

374 @pytest.mark.asyncio 

375 async def test_force_open_missing(self): 

376 reg = CircuitRegistry() 

377 await reg.force_open("nonexistent") # Should not raise 

378 

379 

380# ============================================================================ 

381# Decorator 

382# ============================================================================ 

383 

384class TestCircuitBreakerDecorator: 

385 def test_decorator_creates_breaker(self): 

386 @circuit_breaker("llm_api", failure_threshold=3, timeout_seconds=30) 

387 async def call_llm(prompt: str) -> str: 

388 return f"response: {prompt}" 

389 

390 assert hasattr(call_llm, "_circuit_breaker") 

391 cb = call_llm._circuit_breaker 

392 assert cb.name == "llm_api" 

393 assert cb.config.failure_threshold == 3 

394 

395 @pytest.mark.asyncio 

396 async def test_decorator_success(self): 

397 @circuit_breaker("api") 

398 async def api_call(): 

399 return "ok" 

400 

401 result = await api_call() 

402 assert result == "ok" 

403 

404 @pytest.mark.asyncio 

405 async def test_decorator_failure_trips(self): 

406 @circuit_breaker("api", failure_threshold=1, timeout_seconds=60) 

407 async def failing_call(): 

408 raise ValueError("error") 

409 

410 with pytest.raises(ValueError): 

411 await failing_call() 

412 

413 # Circuit should now be open 

414 with pytest.raises(CircuitOpenError): 

415 await failing_call() 

416 

417 

418# ============================================================================ 

419# CircuitOpenError 

420# ============================================================================ 

421 

422class TestCircuitOpenError: 

423 def test_error_message(self): 

424 err = CircuitOpenError("Circuit 'api' is OPEN") 

425 assert "api" in str(err) 

426 assert "OPEN" in str(err) 

427 

428 

429# ============================================================================ 

430# Default registry 

431# ============================================================================ 

432 

433class TestDefaultRegistry: 

434 def test_default_registry_exists(self): 

435 assert isinstance(default_registry, CircuitRegistry)