Coverage for agentos/tests/test_distributed_lock.py: 100%

1"""Tests for agentos.core.distributed_lock — DistributedLock, InMemoryLockBackend.""" 

2 

3import asyncio 

4from unittest.mock import AsyncMock, MagicMock 

5import pytest 

6from agentos.core.distributed_lock import ( 

7 AbstractLockBackend, 

8 DistributedLock, 

9 InMemoryLockBackend, 

10 LockAcquireError, 

11 LockBackend, 

12 LockConfig, 

13 LockNotHeldError, 

14 LockToken, 

15 PostgresLockBackend, 

16 RedisLockBackend, 

17 _key_to_int64, 

18 create_lock_backend, 

19) 

20 

21 

22# ============================================================================ 

23# LockBackend enum 

24# ============================================================================ 

25 

26class TestLockBackend: 

27 def test_values(self): 

28 assert LockBackend.IN_MEMORY == "in_memory" 

29 assert LockBackend.POSTGRES == "postgres" 

30 assert LockBackend.REDIS == "redis" 

31 

32 

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

34# LockConfig 

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

36 

37class TestLockConfig: 

38 def test_defaults(self): 

39 cfg = LockConfig() 

40 assert cfg.ttl == 30.0 

41 assert cfg.retry_interval == 0.1 

42 assert cfg.acquire_timeout == 10.0 

43 assert cfg.renew_interval == 0.0 

44 

45 def test_custom(self): 

46 cfg = LockConfig(ttl=10.0, retry_interval=0.5, acquire_timeout=5.0, renew_interval=1.0) 

47 assert cfg.ttl == 10.0 

48 assert cfg.renew_interval == 1.0 

49 

50 

51# ============================================================================ 

52# Errors 

53# ============================================================================ 

54 

55class TestErrors: 

56 def test_lock_acquire_error(self): 

57 err = LockAcquireError("timeout") 

58 assert "timeout" in str(err) 

59 

60 def test_lock_not_held_error(self): 

61 err = LockNotHeldError("not held") 

62 assert "not held" in str(err) 

63 

64 

65# ============================================================================ 

66# _key_to_int64 

67# ============================================================================ 

68 

69class TestKeyToInt64: 

70 def test_deterministic(self): 

71 a = _key_to_int64("abc") 

72 b = _key_to_int64("abc") 

73 assert a == b 

74 

75 def test_different_keys(self): 

76 a = _key_to_int64("a") 

77 b = _key_to_int64("b") 

78 assert a != b 

79 

80 

81# ============================================================================ 

82# InMemoryLockBackend 

83# ============================================================================ 

84 

85class TestInMemoryLockBackend: 

86 @pytest.mark.asyncio 

87 async def test_acquire_success(self): 

88 backend = InMemoryLockBackend() 

89 assert await backend.acquire("key1", "owner1", 30.0) 

90 

91 @pytest.mark.asyncio 

92 async def test_acquire_conflict(self): 

93 backend = InMemoryLockBackend() 

94 await backend.acquire("key1", "owner1", 30.0) 

95 assert not await backend.acquire("key1", "owner2", 30.0) 

96 

97 @pytest.mark.asyncio 

98 async def test_acquire_same_owner_reacquire(self): 

99 backend = InMemoryLockBackend() 

100 await backend.acquire("key1", "owner1", 30.0) 

101 # Same owner can re-acquire (overwrite) 

102 assert await backend.acquire("key1", "owner1", 30.0) 

103 

104 @pytest.mark.asyncio 

105 async def test_release_success(self): 

106 backend = InMemoryLockBackend() 

107 await backend.acquire("key1", "owner1", 30.0) 

108 assert await backend.release("key1", "owner1") 

109 

110 @pytest.mark.asyncio 

111 async def test_release_wrong_owner(self): 

112 backend = InMemoryLockBackend() 

113 await backend.acquire("key1", "owner1", 30.0) 

114 assert not await backend.release("key1", "owner2") 

115 

116 @pytest.mark.asyncio 

117 async def test_release_missing_key(self): 

118 backend = InMemoryLockBackend() 

119 assert not await backend.release("nonexistent", "owner1") 

120 

121 @pytest.mark.asyncio 

122 async def test_is_held_true(self): 

123 backend = InMemoryLockBackend() 

124 await backend.acquire("key1", "owner1", 30.0) 

125 assert await backend.is_held("key1", "owner1") 

126 

127 @pytest.mark.asyncio 

128 async def test_is_held_false_wrong_owner(self): 

129 backend = InMemoryLockBackend() 

130 await backend.acquire("key1", "owner1", 30.0) 

131 assert not await backend.is_held("key1", "owner2") 

132 

133 @pytest.mark.asyncio 

134 async def test_is_held_false_missing(self): 

135 backend = InMemoryLockBackend() 

136 assert not await backend.is_held("missing", "owner1") 

137 

138 @pytest.mark.asyncio 

139 async def test_extend_success(self): 

140 backend = InMemoryLockBackend() 

141 await backend.acquire("key1", "owner1", 5.0) 

142 assert await backend.extend("key1", "owner1", 30.0) 

143 

144 @pytest.mark.asyncio 

145 async def test_extend_wrong_owner(self): 

146 backend = InMemoryLockBackend() 

147 await backend.acquire("key1", "owner1", 10.0) 

148 assert not await backend.extend("key1", "owner2", 30.0) 

149 

150 @pytest.mark.asyncio 

151 async def test_extend_missing(self): 

152 backend = InMemoryLockBackend() 

153 assert not await backend.extend("missing", "owner1", 30.0) 

154 

155 @pytest.mark.asyncio 

156 async def test_extend_expired(self): 

157 backend = InMemoryLockBackend() 

158 await backend.acquire("key1", "owner1", 0.001) 

159 await asyncio.sleep(0.01) 

160 assert not await backend.extend("key1", "owner1", 30.0) 

161 

162 @pytest.mark.asyncio 

163 async def test_is_held_false_expired(self): 

164 backend = InMemoryLockBackend() 

165 await backend.acquire("key1", "owner1", 0.001) 

166 await asyncio.sleep(0.01) 

167 assert not await backend.is_held("key1", "owner1") 

168 

169 @pytest.mark.asyncio 

170 async def test_acquire_after_expiry(self): 

171 backend = InMemoryLockBackend() 

172 await backend.acquire("key1", "owner1", 0.001) 

173 await asyncio.sleep(0.01) 

174 assert await backend.acquire("key1", "owner2", 30.0) 

175 

176 

177# ============================================================================ 

178# DistributedLock 

179# ============================================================================ 

180 

181class TestDistributedLock: 

182 @pytest.mark.asyncio 

183 async def test_acquire_success(self): 

184 backend = InMemoryLockBackend() 

185 lock = DistributedLock(backend) 

186 token = await lock.acquire("job:1") 

187 assert isinstance(token, LockToken) 

188 assert token.key == "job:1" 

189 assert await lock.release(token) 

190 

191 @pytest.mark.asyncio 

192 async def test_acquire_timeout(self): 

193 backend = InMemoryLockBackend() 

194 await backend.acquire("job:1", "other", 60.0) 

195 lock = DistributedLock(backend, LockConfig(acquire_timeout=0.05, retry_interval=0.01)) 

196 with pytest.raises(LockAcquireError): 

197 await lock.acquire("job:1") 

198 

199 @pytest.mark.asyncio 

200 async def test_release_false_wrong_owner(self): 

201 backend = InMemoryLockBackend() 

202 lock = DistributedLock(backend) 

203 token = LockToken(key="x", owner_id="wrong", acquired_at=0, ttl=30, _backend=lock) 

204 assert not await lock.release(token) 

205 

206 @pytest.mark.asyncio 

207 async def test_extend(self): 

208 backend = InMemoryLockBackend() 

209 lock = DistributedLock(backend) 

210 token = await lock.acquire("job:1") 

211 assert await lock.extend(token, ttl=60.0) 

212 

213 @pytest.mark.asyncio 

214 async def test_context_manager(self): 

215 backend = InMemoryLockBackend() 

216 lock = DistributedLock(backend) 

217 async with lock("ctx:1") as token: 

218 assert token.key == "ctx:1" 

219 assert await backend.is_held("ctx:1", token.owner_id) 

220 assert not await backend.is_held("ctx:1", token.owner_id) 

221 

222 @pytest.mark.asyncio 

223 async def test_acquire_retry_succeeds(self): 

224 backend = InMemoryLockBackend() 

225 lock = DistributedLock(backend, LockConfig(acquire_timeout=1.0, retry_interval=0.01)) 

226 

227 # Hold lock briefly, then release 

228 async def hold_and_release(): 

229 held_token = await lock.acquire("job:1") 

230 await asyncio.sleep(0.05) 

231 await lock.release(held_token) 

232 

233 task = asyncio.create_task(hold_and_release()) 

234 await asyncio.sleep(0.02) 

235 

236 # This should retry and succeed 

237 token = await lock.acquire("job:1") 

238 assert token.key == "job:1" 

239 await lock.release(token) 

240 await task 

241 

242 

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

244# DistributedLock — Auto-renew 

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

246 

247class TestDistributedLockRenew: 

248 @pytest.mark.asyncio 

249 async def test_auto_renew_starts(self): 

250 backend = InMemoryLockBackend() 

251 lock = DistributedLock(backend, LockConfig(ttl=30, renew_interval=0.05)) 

252 token = await lock.acquire("job:1") 

253 assert token.key in lock._renew_tasks 

254 await lock.release(token) 

255 assert token.key not in lock._renew_tasks 

256 

257 

258# ============================================================================ 

259# create_lock_backend 

260# ============================================================================ 

261 

262class TestCreateLockBackend: 

263 def test_in_memory(self): 

264 backend = create_lock_backend(LockBackend.IN_MEMORY) 

265 assert isinstance(backend, InMemoryLockBackend) 

266 

267 def test_postgres(self): 

268 backend = create_lock_backend(LockBackend.POSTGRES, pool=None) 

269 assert isinstance(backend, PostgresLockBackend) 

270 

271 def test_redis(self): 

272 backend = create_lock_backend(LockBackend.REDIS, redis_client=MagicMock()) 

273 assert isinstance(backend, RedisLockBackend) 

274 

275 def test_redis_missing_client(self): 

276 with pytest.raises(ValueError, match="redis_client"): 

277 create_lock_backend(LockBackend.REDIS) 

278 

279 def test_unknown_backend(self): 

280 with pytest.raises(ValueError, match="Unknown"): 

281 create_lock_backend("fake") # type: ignore 

282 

283 

284# ============================================================================ 

285# PostgresLockBackend (mock) 

286# ============================================================================ 

287 

288class TestPostgresLockBackend: 

289 @pytest.mark.asyncio 

290 async def test_acquire_success(self): 

291 mock_pool = MagicMock() 

292 mock_conn = AsyncMock() 

293 mock_conn.fetchval = AsyncMock(return_value=True) 

294 mock_pool.acquire = AsyncMock(return_value=mock_conn) 

295 mock_pool.release = AsyncMock() 

296 

297 backend = PostgresLockBackend(pool=mock_pool) 

298 assert await backend.acquire("key1", "owner1", 30.0) 

299 

300 @pytest.mark.asyncio 

301 async def test_acquire_failure(self): 

302 mock_pool = MagicMock() 

303 mock_conn = AsyncMock() 

304 mock_conn.fetchval = AsyncMock(return_value=False) 

305 mock_pool.acquire = AsyncMock(return_value=mock_conn) 

306 mock_pool.release = AsyncMock() 

307 

308 backend = PostgresLockBackend(pool=mock_pool) 

309 assert not await backend.acquire("key1", "owner1", 30.0) 

310 

311 @pytest.mark.asyncio 

312 async def test_release_not_acquired(self): 

313 backend = PostgresLockBackend(pool=MagicMock()) 

314 assert not await backend.release("key1", "owner1") 

315 

316 @pytest.mark.asyncio 

317 async def test_release_success(self): 

318 mock_pool = MagicMock() 

319 mock_conn = AsyncMock() 

320 mock_conn.fetchval = AsyncMock(return_value=True) 

321 mock_pool.acquire = AsyncMock(return_value=mock_conn) 

322 mock_pool.release = AsyncMock() 

323 

324 backend = PostgresLockBackend(pool=mock_pool) 

325 await backend.acquire("key1", "owner1", 30.0) 

326 mock_conn.fetchval = AsyncMock(return_value=True) 

327 assert await backend.release("key1", "owner1") 

328 

329 @pytest.mark.asyncio 

330 async def test_extend_held(self): 

331 backend = PostgresLockBackend(pool=MagicMock()) 

332 backend._acquired.add(("key1", "owner1")) 

333 assert await backend.extend("key1", "owner1", 30.0) 

334 

335 @pytest.mark.asyncio 

336 async def test_extend_not_held(self): 

337 backend = PostgresLockBackend(pool=MagicMock()) 

338 assert not await backend.extend("key1", "owner1", 30.0) 

339 

340 @pytest.mark.asyncio 

341 async def test_is_held_true(self): 

342 backend = PostgresLockBackend(pool=MagicMock()) 

343 backend._acquired.add(("key1", "owner1")) 

344 assert await backend.is_held("key1", "owner1") 

345 

346 @pytest.mark.asyncio 

347 async def test_is_held_false(self): 

348 backend = PostgresLockBackend(pool=MagicMock()) 

349 assert not await backend.is_held("key1", "owner1") 

350 

351 @pytest.mark.asyncio 

352 async def test_no_pool_raises(self): 

353 backend = PostgresLockBackend(pool=None) 

354 with pytest.raises(RuntimeError, match="pool"): 

355 await backend.acquire("key1", "owner1", 30.0) 

356 

357 

358# ============================================================================ 

359# RedisLockBackend (mock) 

360# ============================================================================ 

361 

362class TestRedisLockBackend: 

363 @pytest.mark.asyncio 

364 async def test_acquire_success(self): 

365 mock_redis = AsyncMock() 

366 mock_redis.set = AsyncMock(return_value=True) 

367 backend = RedisLockBackend(mock_redis) 

368 assert await backend.acquire("key1", "owner1", 30.0) 

369 

370 @pytest.mark.asyncio 

371 async def test_acquire_failure(self): 

372 mock_redis = AsyncMock() 

373 mock_redis.set = AsyncMock(return_value=None) 

374 backend = RedisLockBackend(mock_redis) 

375 assert not await backend.acquire("key1", "owner1", 30.0) 

376 

377 @pytest.mark.asyncio 

378 async def test_release_fallback_wrong_owner(self): 

379 mock_redis = AsyncMock(spec=["get", "delete"]) 

380 mock_redis.get = AsyncMock(return_value=b"other") 

381 mock_redis.delete = AsyncMock(return_value=True) 

382 backend = RedisLockBackend(mock_redis) 

383 # Fallback path: bytes value is truthy → attempts delete regardless of owner 

384 assert await backend.release("key1", "owner1") 

385 

386 @pytest.mark.asyncio 

387 async def test_release_fallback_success(self): 

388 mock_redis = AsyncMock(spec=["get", "delete"]) 

389 mock_redis.get = AsyncMock(return_value=b"owner1") 

390 mock_redis.delete = AsyncMock(return_value=1) 

391 backend = RedisLockBackend(mock_redis) 

392 assert await backend.release("key1", "owner1") 

393 

394 @pytest.mark.asyncio 

395 async def test_extend_match(self): 

396 mock_redis = AsyncMock() 

397 mock_redis.get = AsyncMock(return_value=b"owner1") 

398 mock_redis.pexpire = AsyncMock(return_value=1) 

399 backend = RedisLockBackend(mock_redis) 

400 assert await backend.extend("key1", "owner1", 30.0) 

401 

402 @pytest.mark.asyncio 

403 async def test_extend_mismatch(self): 

404 mock_redis = AsyncMock() 

405 mock_redis.get = AsyncMock(return_value=b"other") 

406 backend = RedisLockBackend(mock_redis) 

407 assert not await backend.extend("key1", "owner1", 30.0) 

408 

409 @pytest.mark.asyncio 

410 async def test_is_held_true(self): 

411 mock_redis = AsyncMock() 

412 mock_redis.get = AsyncMock(return_value=b"owner1") 

413 backend = RedisLockBackend(mock_redis) 

414 assert await backend.is_held("key1", "owner1") 

415 

416 @pytest.mark.asyncio 

417 async def test_is_held_false(self): 

418 mock_redis = AsyncMock() 

419 mock_redis.get = AsyncMock(return_value=b"other") 

420 backend = RedisLockBackend(mock_redis) 

421 assert not await backend.is_held("key1", "owner1")