482 lines
17 KiB
Python
482 lines
17 KiB
Python
"""
|
|
Runtime bootstrap guard tests.
|
|
|
|
這組測試鎖住 production rollout 曾踩到的兩個啟動序問題：
|
|
- API replicas 同時執行 DB bootstrap DDL 時必須有 advisory lock。
|
|
- SignalWorker 建立 Redis Stream 背景 task 前，必須先初始化 worker Redis pool。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
from collections.abc import Awaitable
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import yaml
|
|
from sqlalchemy.exc import DBAPIError
|
|
|
|
|
|
class _FakeScalarResult:
|
|
def __init__(self, value: bool) -> None:
|
|
self.value = value
|
|
|
|
def scalar(self) -> bool:
|
|
return self.value
|
|
|
|
|
|
class _FakeLockConnection:
|
|
def __init__(self, lock_results: list[bool] | None = None) -> None:
|
|
self.statements: list[str] = []
|
|
self._lock_results = lock_results or [True]
|
|
|
|
async def __aenter__(self) -> _FakeLockConnection:
|
|
return self
|
|
|
|
async def __aexit__(self, *_exc: object) -> None:
|
|
return None
|
|
|
|
async def commit(self) -> None:
|
|
self.statements.append("COMMIT")
|
|
|
|
async def execute(
|
|
self,
|
|
statement: object,
|
|
params: dict[str, str] | None = None,
|
|
) -> _FakeScalarResult:
|
|
sql = str(statement)
|
|
self.statements.append(sql)
|
|
assert params == {"lock_name": "awoooi:init_db:ddl"}
|
|
if "pg_try_advisory_lock" in sql:
|
|
return _FakeScalarResult(self._lock_results.pop(0))
|
|
return _FakeScalarResult(True)
|
|
|
|
|
|
class _FakeEngine:
|
|
def __init__(self, lock_results: list[bool] | None = None) -> None:
|
|
self.lock_conn = _FakeLockConnection(lock_results=lock_results)
|
|
|
|
def connect(self) -> _FakeLockConnection:
|
|
return self.lock_conn
|
|
|
|
|
|
class _ConnectionBudgetEngine:
|
|
def connect(self) -> _FakeLockConnection:
|
|
raise RuntimeError('too many connections for role "awoooi"')
|
|
|
|
|
|
def test_production_worker_delegates_bootstrap_ddl_to_cd() -> None:
|
|
from src.workers import signal_worker
|
|
|
|
root = Path(__file__).resolve().parents[3]
|
|
manifest = yaml.safe_load(
|
|
(root / "k8s/awoooi-prod/08-deployment-worker.yaml").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
container = manifest["spec"]["template"]["spec"]["containers"][0]
|
|
env = {item["name"]: item.get("value") for item in container["env"]}
|
|
source = inspect.getsource(signal_worker._main)
|
|
|
|
assert env["DATABASE_BOOTSTRAP_DDL_ENABLED"] == "false"
|
|
assert "if settings.DATABASE_BOOTSTRAP_DDL_ENABLED" in source
|
|
assert 'owner="cd_migration_preflight"' in source
|
|
|
|
|
|
def test_get_engine_uses_database_pool_budget(monkeypatch):
|
|
from src.db import base as db_base
|
|
|
|
captured: dict[str, object] = {}
|
|
fake_engine = object()
|
|
|
|
def fake_create_async_engine(database_url: str, **kwargs: object) -> object:
|
|
captured["database_url"] = database_url
|
|
captured.update(kwargs)
|
|
return fake_engine
|
|
|
|
monkeypatch.setattr(db_base, "_engine", None)
|
|
monkeypatch.setattr(db_base, "_session_factory", None)
|
|
monkeypatch.setattr(db_base.settings, "DATABASE_URL", "postgresql+asyncpg://u:p@localhost/db")
|
|
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_SIZE", 1)
|
|
monkeypatch.setattr(db_base.settings, "DATABASE_MAX_OVERFLOW", 0)
|
|
monkeypatch.setattr(db_base.settings, "DATABASE_POOL_TIMEOUT_SECONDS", 5.0)
|
|
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", False)
|
|
monkeypatch.setattr(db_base, "create_async_engine", fake_create_async_engine)
|
|
|
|
assert db_base.get_engine() is fake_engine
|
|
assert captured["database_url"] == "postgresql+asyncpg://u:p@localhost/db"
|
|
assert captured["pool_size"] == 1
|
|
assert captured["max_overflow"] == 0
|
|
assert captured["pool_timeout"] == 5.0
|
|
|
|
|
|
def test_get_engine_uses_null_pool_for_bounded_background_processes(monkeypatch):
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from src.db import base as db_base
|
|
|
|
captured: dict[str, object] = {}
|
|
fake_engine = object()
|
|
|
|
def fake_create_async_engine(database_url: str, **kwargs: object) -> object:
|
|
captured["database_url"] = database_url
|
|
captured.update(kwargs)
|
|
return fake_engine
|
|
|
|
monkeypatch.setattr(db_base, "_engine", None)
|
|
monkeypatch.setattr(db_base, "_session_factory", None)
|
|
monkeypatch.setattr(
|
|
db_base.settings,
|
|
"DATABASE_URL",
|
|
"postgresql+asyncpg://u:p@localhost/db",
|
|
)
|
|
monkeypatch.setattr(db_base.settings, "DATABASE_NULL_POOL", True)
|
|
monkeypatch.setattr(db_base, "create_async_engine", fake_create_async_engine)
|
|
|
|
assert db_base.get_engine() is fake_engine
|
|
assert captured["poolclass"] is NullPool
|
|
assert "pool_size" not in captured
|
|
assert "max_overflow" not in captured
|
|
assert "pool_timeout" not in captured
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_db_serializes_bootstrap_ddl_with_advisory_lock(monkeypatch):
|
|
from src.db import base as db_base
|
|
|
|
fake_engine = _FakeEngine()
|
|
calls: list[object] = []
|
|
|
|
async def fake_run_init_db_ddl(engine: object) -> None:
|
|
calls.append(engine)
|
|
|
|
monkeypatch.setattr(db_base, "get_engine", lambda: fake_engine)
|
|
monkeypatch.setattr(db_base, "_run_init_db_ddl", fake_run_init_db_ddl)
|
|
|
|
await db_base.init_db()
|
|
|
|
assert calls == [fake_engine.lock_conn]
|
|
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
|
assert any("pg_advisory_unlock" in stmt for stmt in fake_engine.lock_conn.statements)
|
|
assert "COMMIT" in fake_engine.lock_conn.statements
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_db_skips_bootstrap_when_connection_budget_exhausted(monkeypatch):
|
|
from src.db import base as db_base
|
|
|
|
calls: list[object] = []
|
|
|
|
async def fake_run_init_db_ddl(engine: object) -> None:
|
|
calls.append(engine)
|
|
|
|
monkeypatch.setattr(db_base, "get_engine", lambda: _ConnectionBudgetEngine())
|
|
monkeypatch.setattr(db_base, "_run_init_db_ddl", fake_run_init_db_ddl)
|
|
|
|
await db_base.init_db()
|
|
|
|
assert calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_db_releases_bootstrap_lock_when_ddl_fails(monkeypatch):
|
|
from src.db import base as db_base
|
|
|
|
fake_engine = _FakeEngine()
|
|
|
|
async def fake_run_init_db_ddl(_engine: object) -> None:
|
|
raise RuntimeError("ddl failed")
|
|
|
|
monkeypatch.setattr(db_base, "get_engine", lambda: fake_engine)
|
|
monkeypatch.setattr(db_base, "_run_init_db_ddl", fake_run_init_db_ddl)
|
|
|
|
with pytest.raises(RuntimeError, match="ddl failed"):
|
|
await db_base.init_db()
|
|
|
|
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
|
assert any("pg_advisory_unlock" in stmt for stmt in fake_engine.lock_conn.statements)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_db_skips_bootstrap_when_advisory_lock_times_out(monkeypatch):
|
|
from src.db import base as db_base
|
|
|
|
fake_engine = _FakeEngine(lock_results=[False])
|
|
calls: list[object] = []
|
|
|
|
async def fake_run_init_db_ddl(engine: object) -> None:
|
|
calls.append(engine)
|
|
|
|
monkeypatch.setattr(db_base, "get_engine", lambda: fake_engine)
|
|
monkeypatch.setattr(db_base, "_run_init_db_ddl", fake_run_init_db_ddl)
|
|
monkeypatch.setattr(db_base, "_DB_BOOTSTRAP_LOCK_WAIT_SECONDS", 0.0)
|
|
|
|
await db_base.init_db()
|
|
|
|
assert calls == []
|
|
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
|
assert all("pg_advisory_unlock" not in stmt for stmt in fake_engine.lock_conn.statements)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_db_releases_bootstrap_lock_when_ddl_times_out(monkeypatch):
|
|
from src.db import base as db_base
|
|
|
|
fake_engine = _FakeEngine()
|
|
|
|
async def fake_run_init_db_ddl(_engine: object) -> None:
|
|
raise AssertionError("wait_for should own ddl execution in this test")
|
|
|
|
async def fake_wait_for(coro: Awaitable[Any], timeout: float) -> None:
|
|
assert timeout == 120.0
|
|
coro.close()
|
|
raise TimeoutError
|
|
|
|
monkeypatch.setattr(db_base, "get_engine", lambda: fake_engine)
|
|
monkeypatch.setattr(db_base, "_run_init_db_ddl", fake_run_init_db_ddl)
|
|
monkeypatch.setattr(db_base.asyncio, "wait_for", fake_wait_for)
|
|
|
|
await db_base.init_db()
|
|
|
|
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
|
assert any("pg_advisory_unlock" in stmt for stmt in fake_engine.lock_conn.statements)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_db_skips_bootstrap_when_postgres_statement_timeout(monkeypatch):
|
|
from src.db import base as db_base
|
|
|
|
fake_engine = _FakeEngine()
|
|
|
|
async def fake_run_init_db_ddl(_engine: object) -> None:
|
|
raise DBAPIError(
|
|
"ALTER TABLE approval_records ADD COLUMN IF NOT EXISTS telegram_message_id",
|
|
{},
|
|
Exception("canceling statement due to statement timeout"),
|
|
)
|
|
|
|
monkeypatch.setattr(db_base, "get_engine", lambda: fake_engine)
|
|
monkeypatch.setattr(db_base, "_run_init_db_ddl", fake_run_init_db_ddl)
|
|
|
|
await db_base.init_db()
|
|
|
|
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
|
assert any("pg_advisory_unlock" in stmt for stmt in fake_engine.lock_conn.statements)
|
|
assert "COMMIT" in fake_engine.lock_conn.statements
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_signal_worker_initializes_worker_redis_pool_before_tasks(monkeypatch):
|
|
from src.workers import signal_worker
|
|
|
|
events: list[str] = []
|
|
worker = signal_worker.SignalWorker()
|
|
|
|
async def fake_ensure_consumer_group() -> None:
|
|
events.append("consumer_group")
|
|
|
|
async def fake_init_worker_redis_pool() -> None:
|
|
events.append("worker_redis_pool")
|
|
|
|
def fake_create_task(coro: Awaitable[Any]) -> object:
|
|
events.append("task")
|
|
coro.close()
|
|
|
|
class _FakeTask:
|
|
def cancel(self) -> None:
|
|
return None
|
|
|
|
return _FakeTask()
|
|
|
|
monkeypatch.setattr(worker, "_ensure_consumer_group", fake_ensure_consumer_group)
|
|
monkeypatch.setattr(signal_worker, "init_worker_redis_pool", fake_init_worker_redis_pool)
|
|
monkeypatch.setattr(signal_worker.asyncio, "create_task", fake_create_task)
|
|
|
|
await worker.start()
|
|
|
|
assert events == ["consumer_group", "worker_redis_pool", "task", "task"]
|
|
|
|
|
|
def test_api_lifespan_closes_worker_redis_pool_after_signal_worker() -> None:
|
|
import inspect
|
|
|
|
from src import main as api_main
|
|
|
|
source = inspect.getsource(api_main.lifespan)
|
|
assert "close_worker_redis_pool" in source
|
|
assert source.index("close_signal_worker") < source.index("close_worker_redis_pool")
|
|
assert source.index("close_worker_redis_pool") < source.index("close_redis_pool")
|
|
|
|
|
|
def test_api_background_tasks_default_off_when_db_pool_has_no_spare_slot(monkeypatch) -> None:
|
|
from src import main as api_main
|
|
|
|
monkeypatch.delenv(api_main.API_BACKGROUND_TASKS_ENV, raising=False)
|
|
monkeypatch.setattr(api_main.settings, "DATABASE_POOL_SIZE", 1)
|
|
monkeypatch.setattr(api_main.settings, "DATABASE_MAX_OVERFLOW", 0)
|
|
|
|
assert api_main.should_start_api_background_tasks() is False
|
|
|
|
|
|
def test_api_background_tasks_default_on_when_db_pool_has_spare_slot(monkeypatch) -> None:
|
|
from src import main as api_main
|
|
|
|
monkeypatch.delenv(api_main.API_BACKGROUND_TASKS_ENV, raising=False)
|
|
monkeypatch.setattr(api_main.settings, "DATABASE_POOL_SIZE", 2)
|
|
monkeypatch.setattr(api_main.settings, "DATABASE_MAX_OVERFLOW", 0)
|
|
|
|
assert api_main.should_start_api_background_tasks() is True
|
|
|
|
|
|
def test_api_background_tasks_env_override_controls_serving_budget(monkeypatch) -> None:
|
|
from src import main as api_main
|
|
|
|
monkeypatch.setattr(api_main.settings, "DATABASE_POOL_SIZE", 1)
|
|
monkeypatch.setattr(api_main.settings, "DATABASE_MAX_OVERFLOW", 0)
|
|
|
|
monkeypatch.setenv(api_main.API_BACKGROUND_TASKS_ENV, "true")
|
|
assert api_main.should_start_api_background_tasks() is True
|
|
|
|
monkeypatch.setenv(api_main.API_BACKGROUND_TASKS_ENV, "false")
|
|
assert api_main.should_start_api_background_tasks() is False
|
|
|
|
|
|
def test_api_lifespan_routes_background_loops_through_serving_budget_gate() -> None:
|
|
import inspect
|
|
|
|
from src import main as api_main
|
|
|
|
source = inspect.getsource(api_main.lifespan)
|
|
assert "schedule_api_background_task(run_incident_analysis_sweeper())" in source
|
|
assert "schedule_api_background_task(run_km_backfill_reconciler_loop())" in source
|
|
assert "signal_worker_skipped" in source
|
|
assert "platform_worker_skipped" in source
|
|
|
|
|
|
def test_init_db_backfills_all_optional_knowledge_read_columns() -> None:
|
|
import inspect
|
|
|
|
from src.db import base as db_base
|
|
|
|
source = inspect.getsource(db_base._run_init_db_ddl)
|
|
assert "ALTER TABLE knowledge_entries" in source
|
|
assert "ADD COLUMN IF NOT EXISTS related_approval_id" in source
|
|
assert "ADD COLUMN IF NOT EXISTS path_type" in source
|
|
assert "ADD COLUMN IF NOT EXISTS symptoms_hash" in source
|
|
assert "CREATE INDEX IF NOT EXISTS ix_knowledge_symptoms_hash" in source
|
|
|
|
|
|
def test_init_db_bootstraps_agent99_retry_schema_before_serving() -> None:
|
|
import inspect
|
|
|
|
from src.db import base as db_base
|
|
|
|
source = inspect.getsource(db_base._run_init_db_ddl)
|
|
column_ddl = (
|
|
"ADD COLUMN IF NOT EXISTS next_attempt_at TIMESTAMP NULL"
|
|
)
|
|
index_ddl = "CREATE INDEX IF NOT EXISTS idx_run_state_retry_due"
|
|
|
|
assert "ALTER TABLE awooop_run_state" in source
|
|
assert column_ddl in source
|
|
assert index_ddl in source
|
|
assert "WHERE state = 'pending' AND next_attempt_at IS NOT NULL" in source
|
|
assert source.index(column_ddl) < source.index("ALTER TABLE approval_records")
|
|
|
|
|
|
def test_agent99_retry_migration_matches_startup_bootstrap() -> None:
|
|
migrations = Path(__file__).resolve().parents[1] / "migrations"
|
|
up = (
|
|
migrations / "agent99_dispatch_next_attempt_at_2026-07-11.sql"
|
|
).read_text(encoding="utf-8")
|
|
down = (
|
|
migrations / "agent99_dispatch_next_attempt_at_2026-07-11_down.sql"
|
|
).read_text(encoding="utf-8")
|
|
|
|
assert "ADD COLUMN IF NOT EXISTS next_attempt_at TIMESTAMP NULL" in up
|
|
assert "CREATE INDEX IF NOT EXISTS idx_run_state_retry_due" in up
|
|
assert "WHERE state = 'pending' AND next_attempt_at IS NOT NULL" in up
|
|
assert down.index("DROP INDEX IF EXISTS idx_run_state_retry_due") < (
|
|
down.index("DROP COLUMN IF EXISTS next_attempt_at")
|
|
)
|
|
|
|
|
|
def test_signal_worker_owns_agent99_reconciler_after_schema_bootstrap_gate() -> None:
|
|
import inspect
|
|
|
|
from src import main as api_main
|
|
from src.workers import signal_worker
|
|
|
|
api_source = inspect.getsource(api_main.lifespan)
|
|
worker_source = inspect.getsource(signal_worker._main)
|
|
|
|
assert "run_agent99_controlled_dispatch_reconciler_loop" not in api_source
|
|
assert worker_source.count(
|
|
"run_agent99_controlled_dispatch_reconciler_loop()"
|
|
) == 1
|
|
assert worker_source.index(
|
|
"if settings.DATABASE_BOOTSTRAP_DDL_ENABLED"
|
|
) < worker_source.index(
|
|
"run_agent99_controlled_dispatch_reconciler_loop()"
|
|
)
|
|
assert 'owner="cd_migration_preflight"' in worker_source
|
|
assert 'owner="signal_worker"' in worker_source
|
|
|
|
|
|
def test_signal_worker_agent99_reconciler_requires_url_and_token() -> None:
|
|
import inspect
|
|
|
|
from src.workers import signal_worker
|
|
|
|
source = inspect.getsource(signal_worker._main)
|
|
start = source.index("settings.AGENT99_SRE_ALERT_RELAY_URL.strip()")
|
|
stop = source.index("# Setup graceful shutdown")
|
|
ownership_block = source[start:stop]
|
|
|
|
assert "settings.AGENT99_SRE_ALERT_RELAY_TOKEN.strip()" in ownership_block
|
|
assert "authenticated_relay_not_configured" in ownership_block
|
|
assert (
|
|
"signal_worker_agent99_controlled_dispatch_reconciler_skipped_fail_closed"
|
|
in ownership_block
|
|
)
|
|
|
|
|
|
def test_signal_worker_gracefully_cancels_agent99_reconciler() -> None:
|
|
import inspect
|
|
|
|
from src.workers import signal_worker
|
|
|
|
source = inspect.getsource(signal_worker._main)
|
|
cancel = "agent99_controlled_dispatch_reconciler_task.cancel()"
|
|
await_task = "await agent99_controlled_dispatch_reconciler_task"
|
|
|
|
assert cancel in source
|
|
assert await_task in source
|
|
assert source.index(cancel) < source.index(await_task)
|
|
assert source.index(await_task) < source.index("await worker.stop()")
|
|
|
|
|
|
def test_production_signal_worker_is_single_agent99_reconciler_owner() -> None:
|
|
import inspect
|
|
|
|
from src import main as api_main
|
|
from src.workers import signal_worker
|
|
|
|
root = Path(__file__).resolve().parents[3]
|
|
manifest = yaml.safe_load(
|
|
(root / "k8s/awoooi-prod/08-deployment-worker.yaml").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
container = manifest["spec"]["template"]["spec"]["containers"][0]
|
|
|
|
assert manifest["spec"]["replicas"] == 1
|
|
assert container["command"] == ["python", "-m", "src.workers.signal_worker"]
|
|
assert "run_agent99_controlled_dispatch_reconciler_loop" not in (
|
|
inspect.getsource(api_main.lifespan)
|
|
)
|
|
assert "run_agent99_controlled_dispatch_reconciler_loop" in (
|
|
inspect.getsource(signal_worker._main)
|
|
)
|