mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-08-29 03:56:43 +08:00
284 lines
11 KiB
Python
284 lines
11 KiB
Python
"""
|
|
数据库会话生命周期与资源释放测试。
|
|
|
|
会话生成器（FastAPI 依赖注入入口）必须在请求结束时归还连接，close_database 必须
|
|
释放全部引擎——池化之后引擎不再只有一个：除全局同步/异步引擎外，还有按事件循环
|
|
缓存的池化引擎，漏掉任何一类都是连接泄漏。
|
|
"""
|
|
import asyncio
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from sqlalchemy.orm import scoped_session, sessionmaker
|
|
|
|
from app.runtime.config import global_vars, settings
|
|
from app.db import engine as engine_module
|
|
from app.db import session as session_module
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _restore_pooled_engines():
|
|
"""
|
|
复原按循环缓存的引擎，避免用例间相互污染。
|
|
"""
|
|
saved = dict(session_module._pooled_async_engines)
|
|
yield
|
|
session_module._pooled_async_engines.clear()
|
|
session_module._pooled_async_engines.update(saved)
|
|
|
|
|
|
def test_get_db_closes_session_on_exit(monkeypatch):
|
|
"""
|
|
同步会话生成器必须在迭代结束后关闭会话，否则连接不会归还连接池。
|
|
"""
|
|
closed = []
|
|
fake = MagicMock()
|
|
fake.close = lambda: closed.append(1)
|
|
monkeypatch.setattr(session_module, "SessionFactory", lambda: fake)
|
|
|
|
gen = session_module.get_db()
|
|
assert next(gen) is fake
|
|
with pytest.raises(StopIteration):
|
|
next(gen)
|
|
|
|
assert closed, "生成器正常结束时未关闭会话"
|
|
|
|
|
|
def test_get_db_closes_session_even_on_error(monkeypatch):
|
|
"""
|
|
调用方提前中止时同样要归还会话——否则一次请求失败就泄漏一条连接。
|
|
"""
|
|
closed = []
|
|
fake = MagicMock()
|
|
fake.close = lambda: closed.append(1)
|
|
monkeypatch.setattr(session_module, "SessionFactory", lambda: fake)
|
|
|
|
gen = session_module.get_db()
|
|
next(gen)
|
|
gen.close()
|
|
|
|
assert closed, "生成器被中止时未关闭会话"
|
|
|
|
|
|
def test_get_async_db_yields_session_from_scope(monkeypatch):
|
|
"""
|
|
异步会话入口必须经 async_session_scope 获取——池化与配额都在那里收口，
|
|
绕过它会同时失去连接复用和背压。
|
|
"""
|
|
used = []
|
|
|
|
class _Scope:
|
|
"""会话作用域替身，记录进入与退出。"""
|
|
|
|
async def __aenter__(self):
|
|
used.append("enter")
|
|
return "SESSION"
|
|
|
|
async def __aexit__(self, *_exc):
|
|
used.append("exit")
|
|
return False
|
|
|
|
monkeypatch.setattr(session_module, "async_session_scope", lambda: _Scope())
|
|
|
|
async def run():
|
|
gen = session_module.get_async_db()
|
|
got = await gen.__anext__()
|
|
with pytest.raises(StopAsyncIteration):
|
|
await gen.__anext__()
|
|
return got
|
|
|
|
assert asyncio.run(run()) == "SESSION"
|
|
assert used == ["enter", "exit"], "会话作用域未正确进入/退出"
|
|
|
|
|
|
def test_scoped_sessions_are_not_shared_across_worker_threads(monkeypatch):
|
|
"""同步入口必须为并行工作线程提供不同 Session 实例。"""
|
|
registry = scoped_session(sessionmaker())
|
|
barrier = threading.Barrier(2)
|
|
monkeypatch.setattr(session_module, "_scoped_session", registry)
|
|
|
|
def open_in_thread() -> int:
|
|
"""在线程内持有会话直到另一个线程也完成解析。"""
|
|
session = session_module.ScopedSession()
|
|
try:
|
|
barrier.wait(timeout=5)
|
|
return id(session)
|
|
finally:
|
|
registry.remove()
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
session_ids = list(executor.map(lambda _: open_in_thread(), range(2)))
|
|
|
|
assert len(set(session_ids)) == 2
|
|
|
|
|
|
def test_async_session_scopes_are_not_shared_across_tasks(monkeypatch):
|
|
"""并发异步任务必须各自创建和关闭 AsyncSession 作用域。"""
|
|
created: list[object] = []
|
|
|
|
class FakeAsyncSession:
|
|
"""记录每次作用域构造的独立异步会话替身。"""
|
|
|
|
def __init__(self, **_kwargs) -> None:
|
|
"""创建可由异步上下文管理器返回的唯一实例。"""
|
|
created.append(self)
|
|
|
|
async def __aenter__(self):
|
|
"""返回当前会话实例。"""
|
|
return self
|
|
|
|
async def __aexit__(self, *_exc) -> bool:
|
|
"""模拟正常释放且不吞掉异常。"""
|
|
return False
|
|
|
|
monkeypatch.setattr(
|
|
session_module,
|
|
"_resolve_async_engine",
|
|
lambda: (object(), True),
|
|
)
|
|
monkeypatch.setattr(session_module, "AsyncSession", FakeAsyncSession)
|
|
|
|
async def open_in_task() -> int:
|
|
"""进入一个任务私有的异步会话作用域。"""
|
|
async with session_module.async_session_scope() as session:
|
|
await asyncio.sleep(0)
|
|
return id(session)
|
|
|
|
async def run() -> list[int]:
|
|
"""并发执行两个会话作用域。"""
|
|
return await asyncio.gather(open_in_task(), open_in_task())
|
|
|
|
session_ids = asyncio.run(run())
|
|
|
|
assert len(created) == 2
|
|
assert len(set(session_ids)) == 2
|
|
|
|
|
|
def test_close_database_disposes_pooled_engines(monkeypatch):
|
|
"""
|
|
close_database 必须释放按事件循环缓存的池化引擎。
|
|
|
|
池化之后引擎不再只有全局那一个，漏掉缓存中的引擎意味着进程退出时
|
|
仍持有未归还的物理连接。
|
|
"""
|
|
sync_engine = MagicMock()
|
|
async_engine = MagicMock(dispose=AsyncMock())
|
|
pooled_a = MagicMock(dispose=AsyncMock())
|
|
pooled_b = MagicMock(dispose=AsyncMock())
|
|
|
|
monkeypatch.setattr(engine_module, "_sync_engine", sync_engine)
|
|
monkeypatch.setattr(engine_module, "_async_engine", async_engine)
|
|
session_module._pooled_async_engines.clear()
|
|
session_module._pooled_async_engines.update({1: pooled_a, 2: pooled_b})
|
|
|
|
asyncio.run(session_module.close_database())
|
|
|
|
sync_engine.dispose.assert_called_once()
|
|
async_engine.dispose.assert_awaited_once()
|
|
pooled_a.dispose.assert_awaited_once()
|
|
pooled_b.dispose.assert_awaited_once()
|
|
assert not session_module._pooled_async_engines, "释放后未清空缓存"
|
|
|
|
|
|
def test_close_database_does_not_create_engines_to_dispose_them(monkeypatch):
|
|
"""
|
|
两个引擎槽都是空的时候，close_database 不得为了 dispose 而把引擎创建出来。
|
|
|
|
这是惰性化的直接后果，也是最容易在重构中丢掉的一条：写成 `Engine.dispose()`
|
|
同样能跑通上面那几个用例——它们都把 MagicMock 塞进了引擎槽，`is not None` 恒真，
|
|
于是「先创建再释放」和「有才释放」在测试里完全等价。因此必须单独用空槽压一次：
|
|
否则一个从未用过数据库的进程会在关停时凭空连一次库，只为了随后释放它。
|
|
"""
|
|
created = []
|
|
monkeypatch.setattr(engine_module, "_sync_engine", None)
|
|
monkeypatch.setattr(engine_module, "_async_engine", None)
|
|
monkeypatch.setattr(engine_module, "_get_database_engine",
|
|
lambda **kw: created.append(kw) or MagicMock(dispose=AsyncMock()))
|
|
session_module._pooled_async_engines.clear()
|
|
|
|
asyncio.run(session_module.close_database())
|
|
|
|
assert created == [], f"close_database 为了 dispose 创建了引擎：{created}"
|
|
assert engine_module._sync_engine is None, "同步引擎槽被 close_database 填上了"
|
|
assert engine_module._async_engine is None, "异步引擎槽被 close_database 填上了"
|
|
|
|
|
|
def test_close_database_continues_after_single_engine_failure(monkeypatch):
|
|
"""
|
|
单个引擎释放失败不能中断其余引擎的释放，否则一个坏连接会让其他连接全部泄漏。
|
|
"""
|
|
failing = MagicMock(dispose=AsyncMock(side_effect=RuntimeError("connection reset")))
|
|
healthy = MagicMock(dispose=AsyncMock())
|
|
|
|
monkeypatch.setattr(engine_module, "_sync_engine", MagicMock())
|
|
monkeypatch.setattr(engine_module, "_async_engine", MagicMock(dispose=AsyncMock()))
|
|
session_module._pooled_async_engines.clear()
|
|
session_module._pooled_async_engines.update({1: failing, 2: healthy})
|
|
|
|
asyncio.run(session_module.close_database())
|
|
|
|
healthy.dispose.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.parametrize("failing", ["sync", "async"])
|
|
def test_close_database_releases_remaining_engines_after_global_failure(monkeypatch, failing):
|
|
"""
|
|
全局引擎释放失败：既不能抛出，也不能连累后面的引擎。
|
|
|
|
「不抛」是因为 close_database 在关闭流程末尾调用，抛异常会掩盖其他关闭步骤的问题。
|
|
但只断言「不抛」是不够的——在外面套一个大 try 同样不抛，代价是同步引擎一出错，
|
|
异步引擎和全部池化引擎就都跳过了释放：一条坏连接拖着其余连接一起泄漏，
|
|
而这恰恰是兄弟用例 test_close_database_continues_after_single_engine_failure
|
|
的 docstring 已经声称过的不变量。所以这里把它真正钉住：坏的那个失败，其余照常释放。
|
|
"""
|
|
sync_engine = MagicMock()
|
|
async_engine = MagicMock(dispose=AsyncMock())
|
|
pooled = MagicMock(dispose=AsyncMock())
|
|
if failing == "sync":
|
|
sync_engine.dispose.side_effect = RuntimeError("boom")
|
|
else:
|
|
async_engine.dispose.side_effect = RuntimeError("boom")
|
|
|
|
monkeypatch.setattr(engine_module, "_sync_engine", sync_engine)
|
|
monkeypatch.setattr(engine_module, "_async_engine", async_engine)
|
|
session_module._pooled_async_engines.clear()
|
|
session_module._pooled_async_engines.update({1: pooled})
|
|
|
|
asyncio.run(session_module.close_database()) # 不抛异常
|
|
|
|
# 出错的那个也得真被尝试过，排在它后面的一个都不能少
|
|
sync_engine.dispose.assert_called_once()
|
|
async_engine.dispose.assert_awaited_once()
|
|
pooled.dispose.assert_awaited_once()
|
|
assert not session_module._pooled_async_engines, "释放后未清空缓存"
|
|
|
|
|
|
def test_pooled_engine_is_reused_within_same_loop(monkeypatch):
|
|
"""
|
|
同一事件循环内必须复用同一个池化引擎实例。
|
|
|
|
每次新建引擎等于每次新建一个连接池，连接无法复用，池化就退化回了 NullPool
|
|
的行为——只是多了一层包装。
|
|
"""
|
|
monkeypatch.setattr(settings, "DB_ASYNC_POOL_TYPE", "QueuePool", raising=False)
|
|
created = []
|
|
monkeypatch.setattr(session_module, "_get_database_engine",
|
|
lambda **kw: created.append(kw) or MagicMock())
|
|
|
|
async def run():
|
|
global_vars.CURRENT_EVENT_LOOP = asyncio.get_running_loop()
|
|
session_module._pooled_async_engines.clear()
|
|
first = session_module.get_async_engine()
|
|
second = session_module.get_async_engine()
|
|
return first is second
|
|
|
|
saved = global_vars.CURRENT_EVENT_LOOP
|
|
try:
|
|
assert asyncio.run(run()) is True
|
|
assert len(created) == 1, f"引擎被重复创建 {len(created)} 次"
|
|
assert created[0]["pooled"] is True
|
|
finally:
|
|
global_vars.CURRENT_EVENT_LOOP = saved
|