103 lines
3.2 KiB
Python
103 lines
3.2 KiB
Python
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from app.infra.db import dynamic_manager
|
|
|
|
|
|
class FakeAsyncPool:
|
|
created: list["FakeAsyncPool"] = []
|
|
|
|
def __init__(self, **kwargs) -> None:
|
|
self.kwargs = kwargs
|
|
self.closed = False
|
|
self.created.append(self)
|
|
|
|
async def open(self) -> None:
|
|
return None
|
|
|
|
async def close(self) -> None:
|
|
self.closed = True
|
|
|
|
@asynccontextmanager
|
|
async def connection(self):
|
|
yield object()
|
|
|
|
|
|
def test_active_project_pool_is_not_evicted(monkeypatch) -> None:
|
|
async def exercise() -> None:
|
|
manager = dynamic_manager.ProjectConnectionManager()
|
|
first_id = uuid4()
|
|
second_id = uuid4()
|
|
|
|
async with manager.pg_connection(first_id, "biz_data", "dsn-1", 1, 2):
|
|
first_pool = manager._pg_raw_cache[
|
|
dynamic_manager.CacheKey(first_id, "biz_data")
|
|
].pool
|
|
async with manager.pg_connection(
|
|
second_id, "biz_data", "dsn-2", 1, 2
|
|
):
|
|
assert first_pool.closed is False
|
|
assert len(manager._pg_raw_cache) == 2
|
|
|
|
assert first_pool.closed is False
|
|
assert list(manager._pg_raw_cache) == [
|
|
dynamic_manager.CacheKey(first_id, "biz_data")
|
|
]
|
|
|
|
await manager.close_all()
|
|
|
|
FakeAsyncPool.created = []
|
|
monkeypatch.setattr(dynamic_manager, "AsyncConnectionPool", FakeAsyncPool)
|
|
monkeypatch.setattr(dynamic_manager.settings, "PROJECT_PG_CACHE_SIZE", 1)
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_active_pool_uses_generation_replacement(monkeypatch) -> None:
|
|
async def exercise() -> None:
|
|
manager = dynamic_manager.ProjectConnectionManager()
|
|
project_id = uuid4()
|
|
|
|
async with manager.pg_connection(
|
|
project_id, "biz_data", "old-dsn", 1, 2
|
|
):
|
|
old_pool = FakeAsyncPool.created[0]
|
|
async with manager.pg_connection(
|
|
project_id, "biz_data", "new-dsn", 1, 2
|
|
):
|
|
assert old_pool.closed is False
|
|
assert len(manager._retired_pg) == 1
|
|
|
|
assert old_pool.closed is True
|
|
assert manager._retired_pg == []
|
|
|
|
await manager.close_all()
|
|
|
|
FakeAsyncPool.created = []
|
|
monkeypatch.setattr(dynamic_manager, "AsyncConnectionPool", FakeAsyncPool)
|
|
asyncio.run(exercise())
|
|
|
|
|
|
def test_close_project_does_not_interrupt_active_borrow(monkeypatch) -> None:
|
|
async def exercise() -> None:
|
|
manager = dynamic_manager.ProjectConnectionManager()
|
|
project_id = uuid4()
|
|
key = dynamic_manager.CacheKey(project_id, "iot_data")
|
|
|
|
async with manager.timescale_connection(
|
|
project_id, "iot_data", "ts-dsn", 1, 2
|
|
):
|
|
pool = manager._ts_cache[key].pool
|
|
assert await manager.close_project(project_id) is False
|
|
assert pool.closed is False
|
|
|
|
assert await manager.close_project(project_id) is True
|
|
assert pool.closed is True
|
|
assert key not in manager._ts_cache
|
|
|
|
FakeAsyncPool.created = []
|
|
monkeypatch.setattr(dynamic_manager, "AsyncConnectionPool", FakeAsyncPool)
|
|
asyncio.run(exercise())
|