yuxuanhui
122 lines
5.3 KiB
Python
122 lines
5.3 KiB
Python
"""Stage-three PostgreSQL acceptance in dedicated databases only."""
|
|
|
|
import asyncio
|
|
import os
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from cryptography.fernet import Fernet
|
|
|
|
NAME = "wq_research_stage3_test"
|
|
RESTORE = "wq_research_restore_stage3"
|
|
os.environ.update(
|
|
DATABASE_URL=f"postgresql+asyncpg://postgres:research-test-only@127.0.0.1:18436/{NAME}",
|
|
ADMIN_PASSWORD="research-acceptance-only",
|
|
ENCRYPTION_KEY=Fernet.generate_key().decode(),
|
|
)
|
|
|
|
|
|
def docker(*args, **kwargs):
|
|
return subprocess.run(["docker", "exec", "-i", "wq-research-acceptance-pg", *args], check=True, **kwargs)
|
|
|
|
|
|
async def acceptance():
|
|
from unittest.mock import patch
|
|
|
|
import httpx
|
|
from sqlalchemy import select
|
|
|
|
from app.config import Settings
|
|
from app.main import create_app
|
|
from app.models import TemplateInput
|
|
from app.research.workspace_contracts import TemplateSpec
|
|
from tests.test_ai import configure
|
|
from tests.test_backtests import setup
|
|
from tests.test_research_flows import begin, get, test_fixed_two_rounds_and_idempotent_start
|
|
from tests.test_research_workspace import template
|
|
|
|
app = create_app(Settings(_env_file=None, enable_runner=False, public_origin="http://testserver"))
|
|
async with app.router.lifespan_context(app):
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app), base_url="http://testserver", headers={"X-WQ-Request": "1"}
|
|
) as client:
|
|
assert (
|
|
await client.post(
|
|
"/api/v1/auth/login", json={"username": "admin", "password": "research-acceptance-only"}
|
|
)
|
|
).status_code == 200
|
|
# Configure a deterministic model; no provider or real platform network.
|
|
from tests.ai_fake import fake_model
|
|
|
|
app.state.ai.model_factory = fake_model
|
|
await configure(app, client)
|
|
platform, lane = await setup(app)
|
|
calls = []
|
|
|
|
async def model(ai, context, output_type, revision):
|
|
calls.append(context)
|
|
value = template()
|
|
value["expression"] = f"rank({{field}}) + {len(calls)}"
|
|
return TemplateSpec.model_validate(value), {
|
|
"model": "fixture",
|
|
"revision": revision,
|
|
"usage": {"requests": 1},
|
|
}
|
|
|
|
async with app.state.sessions() as db:
|
|
fixed = await db.scalar(select(TemplateInput))
|
|
body = {
|
|
"request_id": "finite-run",
|
|
"name": "PG 有限研究",
|
|
"input_ids": [fixed.id],
|
|
"hypothesis": "排名稳定性",
|
|
"settings": {"region": "USA", "universe": "TOP3000", "delay": 1},
|
|
"budget": {"max_rounds": 2, "max_simulations": 4, "max_model_calls": 3},
|
|
"batch_candidates": 2,
|
|
}
|
|
with patch("app.research.runtime.request_model", model):
|
|
await test_fixed_two_rounds_and_idempotent_start(app, client, (body, platform, lane, calls))
|
|
short = {
|
|
**body,
|
|
"request_id": "concurrent-budget",
|
|
"budget": {"max_rounds": 1, "max_simulations": 1, "max_model_calls": 2},
|
|
}
|
|
first, second = await asyncio.gather(begin(client, short), begin(client, short))
|
|
assert first["id"] == second["id"]
|
|
from app.research.runtime import ResearchRuntime
|
|
|
|
another = ResearchRuntime(app.state.sessions, app.state.ai, app.state.runner)
|
|
for _ in range(7):
|
|
await asyncio.gather(
|
|
app.state.research.advance(first["id"]), another.advance(first["id"])
|
|
)
|
|
result = await get(client, first["id"])
|
|
assert result["status"] == "budget_exhausted" and result["simulations_used"] == 0
|
|
assert result["model_calls_used"] == 1
|
|
print("PASS PostgreSQL: two-round execution, idempotent starts, concurrent reservations and budget gate")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
docker("createdb", "-U", "postgres", NAME)
|
|
with Path("/tmp/wq-research-stage2.dump").open("rb") as source:
|
|
docker("pg_restore", "-U", "postgres", "-d", NAME, stdin=source)
|
|
config = Config("alembic.ini")
|
|
command.upgrade(config, "0008")
|
|
command.check(config)
|
|
asyncio.run(acceptance())
|
|
dump = Path("/tmp/wq-research-stage3.dump")
|
|
with dump.open("wb") as output:
|
|
docker("pg_dump", "-U", "postgres", "-Fc", NAME, stdout=output)
|
|
docker("createdb", "-U", "postgres", RESTORE)
|
|
with dump.open("rb") as source:
|
|
docker("pg_restore", "-U", "postgres", "-d", RESTORE, stdin=source)
|
|
query = "SELECT (SELECT count(*) FROM research_revisions),(SELECT count(*) FROM research_flow_runs),(SELECT count(*) FROM research_step_runs),(SELECT note FROM research WHERE alpha_id='OLD_RESEARCH')"
|
|
a = docker("psql", "-U", "postgres", "-d", NAME, "-Atc", query, capture_output=True).stdout
|
|
b = docker("psql", "-U", "postgres", "-d", RESTORE, "-Atc", query, capture_output=True).stdout
|
|
assert a == b
|
|
print(
|
|
"PASS PostgreSQL 17: 0007 → 0008 and pg_dump/pg_restore preserve flow budgets, steps and old research notes"
|
|
)
|