feat: add versioned research templates and alpha variants
This commit is contained in:
yuxuanhui
@@ -0,0 +1,203 @@
|
||||
"""Isolated PostgreSQL stage-one migration/restore acceptance; fixed test URL 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
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
URL = "postgresql+asyncpg://postgres:research-test-only@127.0.0.1:18436/wq_research_stage1_test"
|
||||
os.environ.update(
|
||||
DATABASE_URL=URL, ADMIN_PASSWORD="research-acceptance-only", ENCRYPTION_KEY=Fernet.generate_key().decode()
|
||||
)
|
||||
|
||||
|
||||
async def sql(statement):
|
||||
engine = create_async_engine(URL)
|
||||
async with engine.begin() as db:
|
||||
result = await db.execute(text(statement))
|
||||
rows = result.fetchall() if result.returns_rows else None
|
||||
await engine.dispose()
|
||||
return rows
|
||||
|
||||
|
||||
async def acceptance():
|
||||
import httpx
|
||||
|
||||
from app.config import Settings
|
||||
from app.main import create_app
|
||||
from app.worldquant import WqClient
|
||||
from tests.catalog_fake import catalog_response
|
||||
from tests.research_metadata_fake import response as metadata_response
|
||||
from tests.test_backtests import execute, setup, start
|
||||
from tests.test_catalog import prepare, sync
|
||||
from tests.test_research_workspace import expansion, template
|
||||
|
||||
def upstream(request):
|
||||
if request.url.path == "/authentication":
|
||||
return httpx.Response(201, json={})
|
||||
if request.url.path == "/users/self":
|
||||
return httpx.Response(200, json={"id": "PG_RESEARCH_USER"})
|
||||
return metadata_response(request) or catalog_response(request) or httpx.Response(404)
|
||||
|
||||
settings = Settings(_env_file=None, enable_runner=False, public_origin="http://testserver")
|
||||
app = create_app(settings, WqClient(settings, transport=httpx.MockTransport(upstream)))
|
||||
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
|
||||
await client.put(
|
||||
"/api/v1/account/credentials",
|
||||
json={"email": "test@example.com", "password": "synthetic-only"},
|
||||
)
|
||||
job = (await client.post("/api/v1/account/connect")).json()
|
||||
await app.state.runner.execute(job["id"])
|
||||
catalog = (client, app.state.runner, {})
|
||||
await sync(catalog)
|
||||
version = (await sync(catalog, "TEST_FIN"))["id"]
|
||||
fixed = (await prepare(client, version)).json()
|
||||
assert (await client.post("/api/v1/catalog/operators/refresh")).status_code == 200
|
||||
assert (await client.post("/api/v1/catalog/setting-options/refresh")).status_code == 200
|
||||
saved = (
|
||||
await client.post("/api/v1/research/assets", json={"kind": "template", "content": template()})
|
||||
).json()
|
||||
a, b = await asyncio.gather(
|
||||
*(
|
||||
client.put(
|
||||
"/api/v1/research/assets/" + saved["id"],
|
||||
json={"kind": "template", "version": 1, "content": {**template(), "name": name}},
|
||||
)
|
||||
for name in ["A", "B"]
|
||||
)
|
||||
)
|
||||
assert sorted([a.status_code, b.status_code]) == [200, 409]
|
||||
old = (await client.get("/api/v1/research/assets/" + saved["id"] + "?version=1")).json()
|
||||
assert old["name"] == template()["name"]
|
||||
body = expansion(fixed["id"], asset_id=saved["id"], version=1)
|
||||
body.pop("template")
|
||||
experiment = (await client.post("/api/v1/research/experiments", json=body)).json()
|
||||
assert len(experiment["candidates"]) == 2
|
||||
fake, lane = await setup(app)
|
||||
preview = (
|
||||
await client.post("/api/v1/research/experiments/" + experiment["id"] + "/preview", json={})
|
||||
).json()
|
||||
first, second = await asyncio.gather(
|
||||
start(client, preview, "research-confirm"), start(client, preview, "research-confirm")
|
||||
)
|
||||
assert first["backtest_run_id"] == second["backtest_run_id"]
|
||||
await execute(app, lane, first["backtest_run_id"])
|
||||
current = (await client.get("/api/v1/research/experiments/" + experiment["id"])).json()
|
||||
assert current["backtest_run_ids"] == [first["backtest_run_id"]]
|
||||
print(
|
||||
"PASS PostgreSQL: versions, concurrent CAS, fixed input → experiment → idempotent backtest → provenance"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
config = Config("alembic.ini")
|
||||
if asyncio.run(sql("SELECT tablename FROM pg_tables WHERE schemaname='public'")):
|
||||
raise RuntimeError("Dedicated acceptance database must be empty")
|
||||
command.upgrade(config, "0005")
|
||||
asyncio.run(
|
||||
sql(
|
||||
"INSERT INTO alphas (id,hidden,settings,is_metrics,os_metrics,checks,synced_at,raw) VALUES ('OLD_RESEARCH',false,'{}','{}','{}','[]',now(),'{}')"
|
||||
)
|
||||
)
|
||||
asyncio.run(
|
||||
sql(
|
||||
"INSERT INTO research (alpha_id,note,tags,favorite,state,updated_at,version) VALUES ('OLD_RESEARCH','preserve note','[]',false,'inbox',now(),7)"
|
||||
)
|
||||
)
|
||||
command.upgrade(config, "head")
|
||||
command.check(config)
|
||||
assert asyncio.run(sql("SELECT note,version FROM research WHERE alpha_id='OLD_RESEARCH'")) == [
|
||||
("preserve note", 7)
|
||||
]
|
||||
asyncio.run(acceptance())
|
||||
dump = Path("/tmp/wq-research-stage1.dump")
|
||||
with dump.open("wb") as output:
|
||||
subprocess.run(
|
||||
[
|
||||
"docker",
|
||||
"exec",
|
||||
"wq-research-acceptance-pg",
|
||||
"pg_dump",
|
||||
"-U",
|
||||
"postgres",
|
||||
"-Fc",
|
||||
"wq_research_stage1_test",
|
||||
],
|
||||
stdout=output,
|
||||
check=True,
|
||||
)
|
||||
subprocess.run(
|
||||
[
|
||||
"docker",
|
||||
"exec",
|
||||
"wq-research-acceptance-pg",
|
||||
"createdb",
|
||||
"-U",
|
||||
"postgres",
|
||||
"wq_research_restore_stage1",
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
with dump.open("rb") as input_file:
|
||||
subprocess.run(
|
||||
[
|
||||
"docker",
|
||||
"exec",
|
||||
"-i",
|
||||
"wq-research-acceptance-pg",
|
||||
"pg_restore",
|
||||
"-U",
|
||||
"postgres",
|
||||
"-d",
|
||||
"wq_research_restore_stage1",
|
||||
],
|
||||
stdin=input_file,
|
||||
check=True,
|
||||
)
|
||||
query = "SELECT (SELECT count(*) FROM research_revisions),(SELECT count(*) FROM research_experiments),(SELECT count(*) FROM backtest_runs),(SELECT note FROM research WHERE alpha_id='OLD_RESEARCH')"
|
||||
original = subprocess.check_output(
|
||||
[
|
||||
"docker",
|
||||
"exec",
|
||||
"wq-research-acceptance-pg",
|
||||
"psql",
|
||||
"-U",
|
||||
"postgres",
|
||||
"-d",
|
||||
"wq_research_stage1_test",
|
||||
"-Atc",
|
||||
query,
|
||||
]
|
||||
)
|
||||
restored = subprocess.check_output(
|
||||
[
|
||||
"docker",
|
||||
"exec",
|
||||
"wq-research-acceptance-pg",
|
||||
"psql",
|
||||
"-U",
|
||||
"postgres",
|
||||
"-d",
|
||||
"wq_research_restore_stage1",
|
||||
"-Atc",
|
||||
query,
|
||||
]
|
||||
)
|
||||
assert original == restored
|
||||
print(
|
||||
"PASS PostgreSQL 17: 0005 → 0006, schema check, old notes preserved, pg_dump/pg_restore artifacts and provenance counts match"
|
||||
)
|
||||
Reference in New Issue
Block a user