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

1"""Tests for agentos.core.resource_manager — ResourceManager, ResourcePool, ManagedResource.""" 

2 

3import asyncio 

4 

5import pytest 

6 

7from agentos.core.resource_manager import ( 

8 AbstractResource, 

9 ManagedResource, 

10 ResourceInfo, 

11 ResourceManager, 

12 ResourcePool, 

13 ResourceState, 

14 ResourceType, 

15 get_resource_manager, 

16) 

17 

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

19# ResourceInfo 

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

21 

22class TestResourceInfo: 

23 def test_defaults(self): 

24 info = ResourceInfo(name="test") 

25 assert info.name == "test" 

26 assert info.resource_type == ResourceType.OTHER 

27 assert info.state == ResourceState.CREATED 

28 assert info.acquired_at == 0.0 

29 

30 def test_custom(self): 

31 info = ResourceInfo(name="db", resource_type=ResourceType.CONNECTION) 

32 assert info.resource_type == ResourceType.CONNECTION 

33 

34 

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

36# ManagedResource 

37# ============================================================================ 

38 

39class FakeConn: 

40 def __init__(self): 

41 self.closed = False 

42 self.healthy = True 

43 

44 async def close(self): 

45 self.closed = True 

46 

47 async def health_check(self): 

48 return self.healthy 

49 

50 async def ping(self): 

51 return True 

52 

53 

54class TestManagedResource: 

55 @pytest.mark.asyncio 

56 async def test_close(self): 

57 conn = FakeConn() 

58 mr = ManagedResource(conn, name="db") 

59 assert mr.info.state == ResourceState.ACTIVE 

60 await mr.close() 

61 assert mr.info.state == ResourceState.CLOSED 

62 assert conn.closed 

63 

64 @pytest.mark.asyncio 

65 async def test_close_idempotent(self): 

66 conn = FakeConn() 

67 mr = ManagedResource(conn, name="db") 

68 await mr.close() 

69 await mr.close() # should not raise 

70 assert mr.info.state == ResourceState.CLOSED 

71 

72 @pytest.mark.asyncio 

73 async def test_raw_access(self): 

74 conn = FakeConn() 

75 mr = ManagedResource(conn, name="db") 

76 assert mr.raw is conn 

77 

78 @pytest.mark.asyncio 

79 async def test_health_check(self): 

80 conn = FakeConn() 

81 mr = ManagedResource(conn, name="db") 

82 assert await mr.health_check() is True 

83 

84 @pytest.mark.asyncio 

85 async def test_health_check_closed(self): 

86 conn = FakeConn() 

87 mr = ManagedResource(conn, name="db") 

88 await mr.close() 

89 assert await mr.health_check() is False 

90 

91 @pytest.mark.asyncio 

92 async def test_close_fn_fallback(self): 

93 closed_flag = False 

94 

95 class NoCloseObj: 

96 pass 

97 

98 async def custom_close(): 

99 nonlocal closed_flag 

100 closed_flag = True 

101 

102 mr = ManagedResource(NoCloseObj(), close_fn=custom_close, name="x") 

103 await mr.close() 

104 assert closed_flag 

105 

106 @pytest.mark.asyncio 

107 async def test_close_fn_sync(self): 

108 closed = [] 

109 

110 class Obj: 

111 def close(self): 

112 closed.append(1) 

113 

114 mr = ManagedResource(Obj(), name="x") 

115 await mr.close() 

116 assert closed == [1] 

117 

118 @pytest.mark.asyncio 

119 async def test_close_error_sets_error_state(self): 

120 class BadObj: 

121 async def close(self): 

122 raise RuntimeError("boom") 

123 

124 mr = ManagedResource(BadObj(), name="x") 

125 with pytest.raises(RuntimeError, match="boom"): 

126 await mr.close() 

127 assert mr.info.state == ResourceState.ERROR 

128 

129 

130# ============================================================================ 

131# ResourcePool 

132# ============================================================================ 

133 

134class TestResourcePool: 

135 @pytest.mark.asyncio 

136 async def test_acquire_and_release(self): 

137 created = 0 

138 

139 def factory(): 

140 nonlocal created 

141 created += 1 

142 return FakeConn() 

143 

144 pool = ResourcePool(factory, max_size=3) 

145 await pool.start() 

146 async with pool.acquire() as conn: 

147 assert isinstance(conn, FakeConn) 

148 assert not conn.closed 

149 await pool.close() 

150 

151 @pytest.mark.asyncio 

152 async def test_max_size_limit(self): 

153 created = 0 

154 

155 def factory(): 

156 nonlocal created 

157 created += 1 

158 return FakeConn() 

159 

160 pool = ResourcePool(factory, max_size=2) 

161 await pool.start() 

162 

163 async with pool.acquire(): 

164 async with pool.acquire(): 

165 # Third acquire should wait 

166 task = asyncio.ensure_future(pool.acquire().__aenter__()) 

167 await asyncio.sleep(0.01) 

168 assert not task.done() 

169 task.cancel() 

170 try: 

171 await task 

172 except asyncio.CancelledError: 

173 pass 

174 await pool.close() 

175 

176 @pytest.mark.asyncio 

177 async def test_min_size_prefill(self): 

178 created = 0 

179 

180 def factory(): 

181 nonlocal created 

182 created += 1 

183 return FakeConn() 

184 

185 pool = ResourcePool(factory, max_size=5, min_size=2) 

186 await pool.start() 

187 assert created == 2 

188 assert pool.stats["available"] == 2 

189 await pool.close() 

190 

191 @pytest.mark.asyncio 

192 async def test_stats(self): 

193 pool = ResourcePool(FakeConn, max_size=5, min_size=1) 

194 await pool.start() 

195 stats = pool.stats 

196 assert stats["name"] == "pool" 

197 assert stats["max_size"] == 5 

198 assert stats["available"] == 1 

199 assert stats["closed"] is False 

200 await pool.close() 

201 

202 @pytest.mark.asyncio 

203 async def test_close_clears(self): 

204 pool = ResourcePool(FakeConn, max_size=5, min_size=2) 

205 await pool.start() 

206 await pool.close() 

207 assert pool.stats["closed"] is True 

208 

209 @pytest.mark.asyncio 

210 async def test_acquire_after_close(self): 

211 pool = ResourcePool(FakeConn, max_size=3) 

212 await pool.start() 

213 await pool.close() 

214 with pytest.raises(RuntimeError, match="closed"): 

215 async with pool.acquire(): 

216 pass 

217 

218 @pytest.mark.asyncio 

219 async def test_returns_to_pool(self): 

220 created = 0 

221 

222 def factory(): 

223 nonlocal created 

224 created += 1 

225 return FakeConn() 

226 

227 pool = ResourcePool(factory, max_size=3) 

228 await pool.start() 

229 async with pool.acquire(): 

230 pass 

231 # After release, available count is back 

232 assert pool.stats["available"] >= 0 

233 await pool.close() 

234 

235 @pytest.mark.asyncio 

236 async def test_health_check(self): 

237 pool = ResourcePool(FakeConn, max_size=3) 

238 await pool.start() 

239 assert await pool.health_check() is True 

240 await pool.close() 

241 assert await pool.health_check() is False 

242 

243 

244# ============================================================================ 

245# ResourceManager 

246# ============================================================================ 

247 

248class TestResourceManager: 

249 @pytest.mark.asyncio 

250 async def test_register_and_get(self): 

251 rm = ResourceManager() 

252 conn = FakeConn() 

253 mr = ManagedResource(conn, name="db") 

254 await rm.register("db", mr) 

255 assert await rm.get("db") is mr 

256 

257 @pytest.mark.asyncio 

258 async def test_unregister(self): 

259 rm = ResourceManager() 

260 mr = ManagedResource(FakeConn(), name="db") 

261 await rm.register("db", mr) 

262 removed = await rm.unregister("db") 

263 assert removed is mr 

264 assert await rm.get("db") is None 

265 

266 @pytest.mark.asyncio 

267 async def test_shutdown_lifo_order(self): 

268 rm = ResourceManager() 

269 order = [] 

270 

271 class TrackingRes(AbstractResource): 

272 def __init__(self, name): 

273 self.name = name 

274 

275 async def close(self): 

276 order.append(self.name) 

277 

278 await rm.register("a", TrackingRes("a")) 

279 await rm.register("b", TrackingRes("b")) 

280 await rm.register("c", TrackingRes("c")) 

281 

282 await rm.shutdown() 

283 assert order == ["c", "b", "a"] 

284 

285 @pytest.mark.asyncio 

286 async def test_shutdown_reports_failures(self): 

287 rm = ResourceManager() 

288 

289 class FailingRes(AbstractResource): 

290 async def close(self): 

291 raise RuntimeError("fail") 

292 

293 await rm.register("bad", FailingRes()) 

294 failures = await rm.shutdown() 

295 assert len(failures) == 1 

296 assert "bad" in failures[0] 

297 

298 @pytest.mark.asyncio 

299 async def test_health_report(self): 

300 rm = ResourceManager() 

301 conn = FakeConn() 

302 mr = ManagedResource(conn, name="db") 

303 await rm.register("db", mr) 

304 report = await rm.health_report() 

305 assert report["db"] is True 

306 

307 @pytest.mark.asyncio 

308 async def test_finalizer(self): 

309 rm = ResourceManager() 

310 finalizer_ran = [] 

311 

312 rm.add_finalizer(lambda: finalizer_ran.append(1)) 

313 rm.add_finalizer(lambda: finalizer_ran.append(2)) 

314 

315 await rm.register("x", ManagedResource(FakeConn(), name="x")) 

316 await rm.shutdown() 

317 assert finalizer_ran == [1, 2] 

318 

319 @pytest.mark.asyncio 

320 async def test_register_after_shutdown(self): 

321 rm = ResourceManager() 

322 await rm.shutdown() 

323 with pytest.raises(RuntimeError, match="shutting down"): 

324 await rm.register("x", ManagedResource(FakeConn(), name="x")) 

325 

326 @pytest.mark.asyncio 

327 async def test_size(self): 

328 rm = ResourceManager() 

329 assert rm.size == 0 

330 await rm.register("x", ManagedResource(FakeConn(), name="x")) 

331 assert rm.size == 1 

332 

333 def test_check_leaks(self): 

334 rm = ResourceManager() 

335 leaks = rm.check_leaks() 

336 assert leaks == [] 

337 

338 @pytest.mark.asyncio 

339 async def test_shutdown_timeout(self): 

340 rm = ResourceManager() 

341 

342 class SlowRes(AbstractResource): 

343 async def close(self): 

344 await asyncio.sleep(99) 

345 

346 await rm.register("slow", SlowRes()) 

347 failures = await rm.shutdown(timeout=0.001) 

348 assert len(failures) == 1 

349 assert "timeout" in failures[0] 

350 

351 def test_singleton(self): 

352 rm1 = get_resource_manager() 

353 rm2 = get_resource_manager() 

354 assert rm1 is rm2