"""启动时等库(db/base.wait_for_database):和 PostgreSQL 容器一起起、库还没就绪时按退避重试、照实打日志,等够了再报错; 密码不对、库不存在这种等也没用的立刻报错;SQLite 不等。""" from __future__ import annotations import logging import pytest from app.db.base import make_engine, wait_for_database class _Dialect: def __init__(self, name: str) -> None: self.name = name class _Conn: async def __aenter__(self): return self async def __aexit__(self, *exc): return False async def execute(self, _statement): return None class _Engine: """connect() 依次抛 errors 里的错,用完了就连上。""" def __init__(self, errors: list[BaseException], dialect: str = "postgresql") -> None: self.errors = list(errors) self.dialect = _Dialect(dialect) self.attempts = 0 def connect(self): self.attempts += 1 if self.errors: raise self.errors.pop(0) return _Conn() class _Clock: def __init__(self) -> None: self.now = 0.0 self.slept: list[float] = [] def __call__(self) -> float: return self.now async def sleep(self, seconds: float) -> None: self.slept.append(seconds) self.now += seconds class _SqlState(Exception): def __init__(self, sqlstate: str) -> None: super().__init__(f"sqlstate {sqlstate}") self.sqlstate = sqlstate async def test_retries_with_backoff_until_the_database_answers(caplog): clock = _Clock() engine = _Engine([ConnectionRefusedError("拒绝连接"), OSError("名字解析不了"), _SqlState("57P03")]) with caplog.at_level(logging.INFO, logger="app.db.base"): await wait_for_database(engine, timeout=60, sleep=clock.sleep, clock=clock) assert engine.attempts == 4 and clock.slept == [1.0, 2.0, 4.0] warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] assert len(warnings) == 3 and all("应用库还没就绪" in w for w in warnings) assert "拒绝连接" in warnings[0] and "最多 60 秒" in warnings[0] assert any("应用库连上了(第 4 次" in r.getMessage() for r in caplog.records) async def test_gives_up_after_the_budget_and_raises_the_last_error(caplog): clock = _Clock() engine = _Engine([ConnectionRefusedError(f"第 {i} 次") for i in range(100)]) with caplog.at_level(logging.INFO, logger="app.db.base"), pytest.raises(ConnectionRefusedError, match="第 5 次"): await wait_for_database(engine, timeout=20, sleep=clock.sleep, clock=clock) # 退避 1、2、4、5、5:等到第 17 秒,再等 5 秒就过了 20 秒,不等了,抛第 6 次(下标 5)的错 assert clock.slept == [1.0, 2.0, 4.0, 5.0, 5.0] assert any("不再等了" in r.getMessage() for r in caplog.records if r.levelno == logging.ERROR) @pytest.mark.parametrize("sqlstate", ["28P01", "28000", "3D000", "42501"]) async def test_errors_waiting_cannot_fix_are_raised_at_once(sqlstate): clock = _Clock() engine = _Engine([_SqlState(sqlstate)]) with pytest.raises(_SqlState): await wait_for_database(engine, timeout=60, sleep=clock.sleep, clock=clock) assert engine.attempts == 1 and clock.slept == [] async def test_sqlstate_on_the_wrapped_cause_counts(): clock = _Clock() wrapper = RuntimeError("SQLAlchemy 包了一层") wrapper.__cause__ = _SqlState("28P01") with pytest.raises(RuntimeError): await wait_for_database(_Engine([wrapper]), timeout=60, sleep=clock.sleep, clock=clock) assert clock.slept == [] async def test_sqlite_does_not_wait(): engine = _Engine([ConnectionRefusedError("不该连")], dialect="sqlite") await wait_for_database(engine, timeout=60) assert engine.attempts == 0 async def test_a_real_refused_connection_is_treated_as_not_ready_yet(): """真的 asyncpg:端口上没有服务时(库容器还没起来)算「还没就绪」,按退避重试,到时间才抛。""" pytest.importorskip("asyncpg") clock = _Clock() engine = make_engine("postgresql+asyncpg://agentlab:x@127.0.0.1:1/agentlab") try: with pytest.raises(OSError): await wait_for_database(engine, timeout=4, sleep=clock.sleep, clock=clock) finally: await engine.dispose() assert clock.slept == [1.0, 2.0] async def test_a_missing_database_on_a_real_server_is_raised_at_once(): """真的 PostgreSQL(AGENTLAB_TEST_DATABASE_URL):库不存在(3D000)立刻报错,不等。""" from app.core.config import settings if settings.db_dialect != "postgresql": pytest.skip("要 PostgreSQL") base, _, _ = settings.db_url.rpartition("/") clock = _Clock() engine = make_engine(f"{base}/agentlab_no_such_database") try: with pytest.raises(Exception) as caught: await wait_for_database(engine, timeout=60, sleep=clock.sleep, clock=clock) finally: await engine.dispose() assert clock.slept == [], caught.value assert "agentlab_no_such_database" in str(caught.value) async def test_startup_waits_for_the_database_before_the_version_check(monkeypatch): """lifespan 里第一次连库的是版本检查(cluster/version.preflight):等库要在它前面,库的容器晚一步起来时服务不至于直接崩。""" from app.cluster import guard, version from app.db import base from app.main import app, lifespan calls: list[str] = [] async def waited(*_a, **_k) -> None: calls.append("wait") async def checked(*_a, **_k) -> None: calls.append("preflight") raise RuntimeError("到这里就够了") monkeypatch.setattr(guard, "acquire", lambda: None) monkeypatch.setattr(base, "wait_for_database", waited) monkeypatch.setattr(version, "preflight", checked) with pytest.raises(RuntimeError, match="到这里就够了"): async with lifespan(app): pass assert calls == ["wait", "preflight"]