98 lines
4.8 KiB
Python
98 lines
4.8 KiB
Python
|
"""Disposable PostgreSQL migration/concurrency acceptance; synthetic platform only.
|
||
|
|
|
||
|
|
MCP_TEST_DATABASE_URL must identify a local database named wq_mcp_test.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import os
|
||
|
|
from urllib.parse import urlsplit
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
from alembic import command
|
||
|
|
from alembic.config import Config
|
||
|
|
from cryptography.fernet import Fernet
|
||
|
|
from sqlalchemy import func, select
|
||
|
|
|
||
|
|
from app.alphas import upsert_alpha
|
||
|
|
from app.config import Settings
|
||
|
|
from app.db import create_database
|
||
|
|
from app.main import create_app
|
||
|
|
from app.mcp_api.auth import SCOPES, authenticate, create_token
|
||
|
|
from app.models import BacktestEvent, BacktestPreview, BacktestRun, Research, ResearchRequest
|
||
|
|
from app.research_access.contracts import Control, Submit
|
||
|
|
from app.research_access.service import ResearchAccess, ResearchError
|
||
|
|
from app.worldquant import WqClient
|
||
|
|
from tests.backtest_fake import Platform
|
||
|
|
from tests.test_backtests import candidate, setup
|
||
|
|
|
||
|
|
|
||
|
|
async def seed(settings):
|
||
|
|
engine, sessions = create_database(settings.database_url)
|
||
|
|
async with sessions.begin() as db:
|
||
|
|
await upsert_alpha(db, {"id": "legacy", "regular": {"code": "rank(close)"}})
|
||
|
|
(await db.get(Research, "legacy")).note = "keep historical note"
|
||
|
|
await engine.dispose()
|
||
|
|
|
||
|
|
|
||
|
|
async def verify(settings):
|
||
|
|
app = create_app(settings, WqClient(settings, transport=httpx.MockTransport(Platform())))
|
||
|
|
async with app.router.lifespan_context(app):
|
||
|
|
await setup(app)
|
||
|
|
async with app.state.sessions.begin() as db:
|
||
|
|
_, secret = await create_token(db, "synthetic", SCOPES)
|
||
|
|
principal = await authenticate(db, secret)
|
||
|
|
args = Submit(name="concurrent", candidates=[candidate()], idempotency_key="same")
|
||
|
|
async def submit(body):
|
||
|
|
async with app.state.sessions.begin() as db:
|
||
|
|
return await ResearchAccess(db, principal, app.state.runner.client, settings.public_origin).submit(body)
|
||
|
|
a, b = await asyncio.gather(submit(args), submit(args))
|
||
|
|
assert a == b
|
||
|
|
rid = a["backtest_run_id"]
|
||
|
|
async with app.state.sessions() as db:
|
||
|
|
assert await db.scalar(select(func.count()).select_from(BacktestRun)) == 1
|
||
|
|
assert await db.scalar(select(func.count()).select_from(BacktestPreview)) == 1
|
||
|
|
assert await db.scalar(select(func.count()).select_from(ResearchRequest)) == 1
|
||
|
|
assert (await db.get(Research, "legacy")).note == "keep historical note"
|
||
|
|
try:
|
||
|
|
await submit(args.model_copy(update={"name": "different"}))
|
||
|
|
except ResearchError as exc:
|
||
|
|
assert exc.data["code"] == "IDEMPOTENCY_CONFLICT"
|
||
|
|
else:
|
||
|
|
raise AssertionError("different payload accepted")
|
||
|
|
control = Control(run_id=rid, action="pause", expected_version=1, idempotency_key="pause")
|
||
|
|
async def pause():
|
||
|
|
async with app.state.sessions.begin() as db:
|
||
|
|
return await ResearchAccess(db, principal, app.state.runner.client, settings.public_origin).control(control)
|
||
|
|
x, y = await asyncio.gather(pause(), pause())
|
||
|
|
assert x == y
|
||
|
|
async with app.state.sessions() as db:
|
||
|
|
assert await db.scalar(select(func.count()).select_from(BacktestEvent).where(BacktestEvent.kind == "control")) == 1
|
||
|
|
# A fresh process-equivalent application returns the original response after restart.
|
||
|
|
replacement = create_app(settings, WqClient(settings, transport=httpx.MockTransport(Platform())))
|
||
|
|
async with replacement.router.lifespan_context(replacement):
|
||
|
|
async with replacement.state.sessions.begin() as db:
|
||
|
|
principal = await authenticate(db, secret)
|
||
|
|
replay = await ResearchAccess(db, principal, replacement.state.runner.client, settings.public_origin).submit(args)
|
||
|
|
assert replay == a
|
||
|
|
print("PostgreSQL: migration preserves history; concurrent submit/control and restart replay passed")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
url = os.environ["MCP_TEST_DATABASE_URL"]
|
||
|
|
parsed = urlsplit(url)
|
||
|
|
if parsed.hostname not in {"127.0.0.1", "localhost"} or parsed.path != "/wq_mcp_test":
|
||
|
|
raise SystemExit("Refusing a non-local/non-disposable database")
|
||
|
|
key = Fernet.generate_key().decode()
|
||
|
|
os.environ.update(DATABASE_URL=url, ADMIN_PASSWORD="synthetic-admin-only", ENCRYPTION_KEY=key,
|
||
|
|
WQ_EMAIL="", WQ_PASSWORD="")
|
||
|
|
settings = Settings(_env_file=None, database_url=url, admin_password="synthetic-admin-only",
|
||
|
|
encryption_key=key, enable_runner=False, mcp_enabled=True, public_origin="http://testserver")
|
||
|
|
config = Config("alembic.ini")
|
||
|
|
command.upgrade(config, "0009")
|
||
|
|
asyncio.run(seed(settings))
|
||
|
|
command.upgrade(config, "head")
|
||
|
|
asyncio.run(verify(settings))
|
||
|
|
command.downgrade(config, "0009")
|
||
|
|
command.upgrade(config, "head")
|
||
|
|
print("PostgreSQL: additive migration downgrade/re-upgrade passed")
|