NeteaseUrl/tests/test_db_conn.py

"""db.connect 连接策略契约：WAL 首连设置一次 + 目录自动创建。

背景：原同步实现每个连接都执行 mkdir + PRAGMA journal_mode=WAL，实测每次
多付 ~1.2ms（mkdir 190us + 新连接上 PRAGMA ~1ms）。WAL 模式持久化在库文件头
里，与 core/db.py 的 _WAL_READY_FOR 同款，按路径设一次即可。迁移到 aiosqlite
后策略不变：`_configure` 在专用锁内首连设置一次，库文件被删重建时（force_wal）
再补设。
"""

import asyncio
import threading
from pathlib import Path

import aiosqlite

from auth_system import db


def _point_db_to(monkeypatch, tmp_path: Path, statements: list[str] | None = None):
    """把 DB_PATH 指到临时库，并清空 WAL 已就绪缓存。"""
    new_path = tmp_path / 'site.db'
    monkeypatch.setattr(db, 'DB_PATH', new_path)
    monkeypatch.setattr(db, '_WAL_READY_FOR', None)
    if statements is not None:
        real_connect = aiosqlite.connect

        async def tracing_connect(*a, **k):
            # 在第一句 PRAGMA 之前挂追踪器，才能捕获 _configure 执行的语句
            conn = await real_connect(*a, **k)
            await conn.set_trace_callback(statements.append)
            return conn

        monkeypatch.setattr(db.aiosqlite, 'connect', tracing_connect)
    return new_path


async def _open_close() -> None:
    conn = await db.connect()
    await conn.close()


async def _journal_mode() -> str:
    conn = await db.connect()
    try:
        row = await (await conn.execute('PRAGMA journal_mode')).fetchone()
        assert row is not None
        return str(row[0]).lower()
    finally:
        await conn.close()


def test_wal_mode_active_on_fresh_db(monkeypatch, tmp_path):
    """新路径首连：WAL 生效。"""
    _point_db_to(monkeypatch, tmp_path)
    assert asyncio.run(_journal_mode()) == 'wal'


def test_wal_pragma_only_once_per_path(monkeypatch, tmp_path):
    """同一路径第二次起不再执行 PRAGMA journal_mode=WAL（SQL 追踪器实证）。"""
    statements: list[str] = []
    _point_db_to(monkeypatch, tmp_path, statements=statements)

    asyncio.run(_open_close())
    assert any('journal_mode=WAL' in s for s in statements), f'首连必须设 WAL，实际: {statements}'

    statements.clear()
    asyncio.run(_open_close())
    assert not any('journal_mode=WAL' in s for s in statements), \
        f'第二条连接不应再设 WAL，实际: {statements}'
    # 常规 PRAGMA（busy_timeout 等）每连接仍照设
    assert any('busy_timeout' in s for s in statements), statements


def test_new_path_resets_wal_cache(monkeypatch, tmp_path):
    """换库路径（测试隔离场景）：新路径重新设置 WAL。"""
    _point_db_to(monkeypatch, tmp_path / 'a')
    asyncio.run(_open_close())

    monkeypatch.setattr(db, 'DB_PATH', tmp_path / 'b' / 'site.db')
    assert asyncio.run(_journal_mode()) == 'wal'


def test_missing_dir_auto_created(monkeypatch, tmp_path):
    """目录不存在：connect 失败兜底 mkdir 重试，不经 init_db 也能用。"""
    deep = tmp_path / 'x' / 'y'
    _point_db_to(monkeypatch, deep)
    assert not deep.exists()

    async def _main():
        conn = await db.connect()
        try:
            await conn.execute('CREATE TABLE t (v)')
            await conn.commit()
        finally:
            await conn.close()

    asyncio.run(_main())
    assert (deep / 'site.db').exists()


def test_deleted_db_file_recovers(monkeypatch, tmp_path):
    """库文件被删（目录还在）：重连自动重建且恢复 WAL 模式，不崩。"""
    _point_db_to(monkeypatch, tmp_path)
    asyncio.run(_open_close())

    (tmp_path / 'site.db').unlink()
    assert asyncio.run(_journal_mode()) == 'wal', '重建的库文件必须恢复 WAL（journal_mode 不持久于新文件）'


def test_concurrent_first_connect_no_lease_conflict(monkeypatch, tmp_path):
    """**单事件循环内**并发首连（生产模型：一个 loop 上并发多个请求）：
    专用锁 + 二次复查保证 WAL 只切一次，无连接报错。
    """
    _point_db_to(monkeypatch, tmp_path)

    async def _main():
        async def _one():
            conn = await db.connect()
            try:
                await (await conn.execute('SELECT 1')).fetchall()
            finally:
                await conn.close()

        await asyncio.gather(*[_one() for _ in range(16)])

    asyncio.run(_main())


def test_concurrent_connect_from_threads_after_wal_ready(monkeypatch, tmp_path):
    """WAL 已就绪后，多线程各自独立 loop 建连不再触排他锁争用（读路径并发）。

    为何不再钉「多线程**同时首连**」：异步层的 WAL 切换锁按事件循环隔离
    （见 auth_system/db.py 的 _loop_lock，与 core/db.py 同款），跨 loop 的首连
    无法用同一把 asyncio.Lock 串行化——那是同步栈 threading.Lock 才有的保证。
    生产只有一个事件循环（任务监控线程的 asyncio.run 只在 init_db 之后启动），
    首连争用实际发生在单 loop 内，由上一条用例覆盖。
    """
    _point_db_to(monkeypatch, tmp_path)
    asyncio.run(_open_close())   # 先把库切到 WAL，置 _WAL_READY_FOR
    errors: list[Exception] = []

    def worker():
        async def _one():
            conn = await db.connect()
            try:
                await (await conn.execute('SELECT 1')).fetchall()
            finally:
                await conn.close()

        try:
            for _ in range(5):
                asyncio.run(_one())
        except Exception as e:  # 收集全部失败供断言
            errors.append(e)

    threads = [threading.Thread(target=worker) for _ in range(8)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    assert errors == []