yuxuanhui
122 lines
5.0 KiB
Python
122 lines
5.0 KiB
Python
"""Isolated PostgreSQL migration/concurrency acceptance. Never point at a personal database.
|
|
|
|
Run with DATABASE_URL ending in /wq_backtest_test, synthetic ADMIN_PASSWORD and
|
|
ENCRYPTION_KEY. Uses only mock WorldQuant HTTP and a disposable database.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
|
|
import httpx
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
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.models import BacktestEvent, BacktestResult, BacktestRun, Research, SimulationAttempt
|
|
from app.worldquant import WqClient
|
|
from tests.backtest_fake import Platform
|
|
from tests.test_backtests import candidate, preview, setup, start, tick
|
|
|
|
|
|
async def seed_old(settings):
|
|
engine, sessions = create_database(settings.database_url)
|
|
async with sessions.begin() as db:
|
|
await upsert_alpha(
|
|
db,
|
|
{
|
|
"id": "MIGRATION_ALPHA",
|
|
"type": "REGULAR",
|
|
"regular": {"code": "rank(close) + 0"},
|
|
"settings": candidate()["settings"],
|
|
},
|
|
)
|
|
await db.flush()
|
|
research = await db.get(Research, "MIGRATION_ALPHA")
|
|
research.note = "keep old research across upgrade and simulations"
|
|
await engine.dispose()
|
|
|
|
|
|
async def acceptance(settings):
|
|
fake = Platform()
|
|
app = create_app(settings, WqClient(settings, transport=httpx.MockTransport(fake)))
|
|
async with app.router.lifespan_context(app):
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=app),
|
|
base_url="http://testserver",
|
|
headers={"X-WQ-Request": "1"},
|
|
) as client:
|
|
assert (
|
|
await client.post(
|
|
"/api/v1/auth/login",
|
|
json={"username": "admin", "password": settings.admin_password.get_secret_value()},
|
|
)
|
|
).status_code == 200
|
|
fake, lane = await setup(app)
|
|
fake.existing_alpha_ids = ["MIGRATION_ALPHA"]
|
|
p = await preview(client, [candidate(0), candidate(0) | {"client_item_id": "repeat"}])
|
|
a, b = await asyncio.gather(
|
|
start(client, p, "concurrent-confirm"), start(client, p, "concurrent-confirm")
|
|
)
|
|
assert a["backtest_run_id"] == b["backtest_run_id"]
|
|
rid = a["backtest_run_id"]
|
|
for _ in range(5):
|
|
await tick(lane)
|
|
result = (await client.get(f"/api/v1/backtests/runs/{rid}")).json()
|
|
assert result["status"] == "completed", result
|
|
assert len(fake.posts) == 2
|
|
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(BacktestResult)) == 2
|
|
note = (await db.get(Research, "MIGRATION_ALPHA")).note
|
|
assert note == "keep old research across upgrade and simulations"
|
|
events = list(
|
|
await db.scalars(
|
|
select(BacktestEvent.seq)
|
|
.where(BacktestEvent.run_id == rid)
|
|
.order_by(BacktestEvent.seq)
|
|
)
|
|
)
|
|
assert events == list(range(1, len(events) + 1))
|
|
# Leave an accepted run for a new application instance to recover.
|
|
next_run = await start(client, await preview(client, [candidate(2)]), "restart")
|
|
async with app.state.sessions() as db:
|
|
aid = await db.scalar(
|
|
select(SimulationAttempt.id).where(
|
|
SimulationAttempt.run_id == next_run["backtest_run_id"]
|
|
)
|
|
)
|
|
await lane.step(aid)
|
|
await lane.interrupt()
|
|
replacement = create_app(settings, WqClient(settings, transport=httpx.MockTransport(fake)))
|
|
async with replacement.router.lifespan_context(replacement):
|
|
lane = replacement.state.runner.backtests
|
|
await lane.start()
|
|
await lane.stop()
|
|
await lane.step(aid)
|
|
async with replacement.state.sessions() as db:
|
|
assert (await db.get(BacktestRun, next_run["backtest_run_id"])).status == "completed"
|
|
assert len(fake.posts) == 3
|
|
print(
|
|
"PASS PostgreSQL: concurrent confirmation creates one run; two attempts share one Alpha safely; contiguous transactional events; research preserved; replacement application resumes accepted simulation without POST"
|
|
)
|
|
|
|
|
|
def main():
|
|
if not os.environ.get("DATABASE_URL", "").endswith("/wq_backtest_test"):
|
|
raise SystemExit("Only an isolated wq_backtest_test database is allowed")
|
|
settings = Settings(_env_file=None, enable_runner=False, public_origin="http://testserver")
|
|
config = Config("alembic.ini")
|
|
command.upgrade(config, "0002")
|
|
asyncio.run(seed_old(settings))
|
|
command.upgrade(config, "head")
|
|
command.check(config)
|
|
asyncio.run(acceptance(settings))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|