65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
"""
|
||
并发启动竞态测试：模拟 gunicorn 多 worker 同时启动时，
|
||
多个进程并发调用 plugins.auth.create_workers 初始化默认管理员。
|
||
|
||
预期：不抛 IntegrityError，users 表中只有一条 admin 记录。
|
||
"""
|
||
import asyncio
|
||
import sys
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).parent / "src"))
|
||
|
||
import plugins.auth as auth_plugin
|
||
|
||
|
||
def fake_create_workers(session, config, count=1, worker_manager=None):
|
||
"""替换真实 worker 启动，避免依赖 zenoh 会话。与原始实现同为同步函数。"""
|
||
return []
|
||
|
||
|
||
async def main():
|
||
auth_plugin._create_workers = fake_create_workers
|
||
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
config = {"type": "auth", "data_dir": tmp}
|
||
|
||
# 模拟 4 个 gunicorn worker 进程同时启动
|
||
results = await asyncio.gather(
|
||
auth_plugin.create_workers(None, config, 3),
|
||
auth_plugin.create_workers(None, config, 3),
|
||
auth_plugin.create_workers(None, config, 3),
|
||
auth_plugin.create_workers(None, config, 3),
|
||
return_exceptions=True,
|
||
)
|
||
|
||
errors = [r for r in results if isinstance(r, Exception)]
|
||
if errors:
|
||
for e in errors:
|
||
print("❌ 异常:", type(e).__name__, e)
|
||
raise RuntimeError("并发启动存在未捕获异常")
|
||
|
||
# 校验 users 表中应只有一条 admin
|
||
from chongming_database import Database
|
||
from plugins.auth.models.auth import Users
|
||
from sqlalchemy import select
|
||
|
||
db = Database(Path(tmp) / "auth.db")
|
||
await db.initialize()
|
||
async with db.get_session() as sess:
|
||
res = await sess.execute(select(Users))
|
||
rows = res.scalars().all()
|
||
|
||
print(f"✅ 4 个并发启动全部成功")
|
||
print(f"✅ users 表记录数: {len(rows)}（应为 1）")
|
||
for u in rows:
|
||
print(f" - {u.username} <{u.email}> superuser={u.is_superuser}")
|
||
|
||
assert len(rows) == 1, f"预期 1 条 admin，实际 {len(rows)} 条"
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|
||
print("🎉 竞态测试通过")
|