88 lines
4.8 KiB
Python
88 lines
4.8 KiB
Python
|
"""Retire template scope once in storage, without a runtime compatibility path."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import sqlalchemy as sa
|
||
|
|
from alembic import command
|
||
|
|
from alembic.config import Config
|
||
|
|
from cryptography.fernet import Fernet
|
||
|
|
from sqlalchemy.orm import Session
|
||
|
|
|
||
|
|
from app.models import (
|
||
|
|
Account,
|
||
|
|
ResearchAsset,
|
||
|
|
ResearchExperiment,
|
||
|
|
ResearchFlowRun,
|
||
|
|
ResearchRequest,
|
||
|
|
ResearchRevision,
|
||
|
|
ResearchStepRun,
|
||
|
|
)
|
||
|
|
from app.research.workspace_contracts import FeatureSpec, TemplateSpec
|
||
|
|
|
||
|
|
|
||
|
|
def test_remove_template_scope_preserves_execution_scopes_and_all_versions(tmp_path, monkeypatch):
|
||
|
|
database = tmp_path / "template-scopes.db"
|
||
|
|
monkeypatch.setenv("DATABASE_URL", f"sqlite+aiosqlite:///{database}")
|
||
|
|
monkeypatch.setenv("ADMIN_PASSWORD", "migration-test-only")
|
||
|
|
monkeypatch.setenv("ENCRYPTION_KEY", Fernet.generate_key().decode())
|
||
|
|
monkeypatch.setenv("WQ_EMAIL", "")
|
||
|
|
monkeypatch.setenv("WQ_PASSWORD", "")
|
||
|
|
root = Path(__file__).resolve().parents[1]
|
||
|
|
config = Config(str(root / "alembic.ini"))
|
||
|
|
config.set_main_option("script_location", str(root / "migrations"))
|
||
|
|
command.upgrade(config, "0022")
|
||
|
|
engine = sa.create_engine(f"sqlite:///{database}")
|
||
|
|
scope = {"instrument_type": "EQUITY", "region": "EUR", "universe": "TOP2500", "delay": 1}
|
||
|
|
clean = {"name": "template", "expression": "rank({scope})", "description": "scope is also a variable name",
|
||
|
|
"variables": {"scope": {"kind": "field", "field_type": "MATRIX", "values": []}}}
|
||
|
|
old = {**clean, "scope": scope}
|
||
|
|
feature = {"name": "feature", "hypothesis": "test", "template": old}
|
||
|
|
clean_feature = {**feature, "template": clean}
|
||
|
|
payload = {"template": {"id": "a", "version": 1, "content": old}, "feature": {"content": feature},
|
||
|
|
"inputs": [{"scope": scope, "field_ids": ["close"]}], "scope": scope,
|
||
|
|
"bindings": {"scope": "close"}, "items": [old, None, "unchanged"]}
|
||
|
|
expected = {**payload, "template": {**payload["template"], "content": clean},
|
||
|
|
"feature": {"content": clean_feature}, "items": [clean, None, "unchanged"]}
|
||
|
|
with Session(engine) as db, db.begin():
|
||
|
|
db.add(Account(id=1))
|
||
|
|
db.add_all([ResearchAsset(id="a", kind="template", name="template", version=501),
|
||
|
|
ResearchAsset(id="b", kind="feature", name="feature", version=2)])
|
||
|
|
db.flush()
|
||
|
|
db.add_all([ResearchRevision(asset_id="a", version=i, content={**old, "scope": scope if i % 2 else None},
|
||
|
|
provenance=payload) for i in range(1, 502)])
|
||
|
|
db.add_all([ResearchRevision(asset_id="b", version=i, content=feature, provenance=payload) for i in (1, 2)])
|
||
|
|
db.add(ResearchExperiment(id="experiment", name="test", kind="template", hypothesis="test",
|
||
|
|
inputs=payload["inputs"], parents=[payload], candidates=[{"settings": scope}], evidence=payload))
|
||
|
|
db.add(ResearchFlowRun(id="flow", request_id="request", name="test", definition=payload, authorization=payload))
|
||
|
|
db.flush()
|
||
|
|
db.add(ResearchStepRun(id="step", run_id="flow", node_id="expand", round=1, output=payload))
|
||
|
|
db.add(ResearchRequest(id="request", account_id=1, operation="create_research_template",
|
||
|
|
idempotency_key="key", digest="a" * 64, business_id="a", response=payload))
|
||
|
|
# Include every stored column in the comparison, not just the removed key.
|
||
|
|
tables = {name: sa.Table(name, sa.MetaData(), autoload_with=engine) for name in (
|
||
|
|
"research_revisions", "research_experiments", "research_flow_runs", "research_step_runs", "research_requests",
|
||
|
|
)}
|
||
|
|
def snapshots():
|
||
|
|
with engine.connect() as db:
|
||
|
|
return {name: [dict(row) for row in db.execute(sa.select(table).order_by(*table.primary_key.columns)).mappings()]
|
||
|
|
for name, table in tables.items()}
|
||
|
|
before = snapshots()
|
||
|
|
wanted = {name: [dict(row) for row in rows] for name, rows in before.items()}
|
||
|
|
for row in wanted["research_revisions"]:
|
||
|
|
row.update(content=clean if row["asset_id"] == "a" else clean_feature, provenance=expected)
|
||
|
|
wanted["research_experiments"][0].update(evidence=expected, parents=[expected])
|
||
|
|
wanted["research_flow_runs"][0].update(definition=expected, authorization=expected)
|
||
|
|
wanted["research_step_runs"][0].update(output=expected)
|
||
|
|
wanted["research_requests"][0].update(response=expected)
|
||
|
|
command.upgrade(config, "head")
|
||
|
|
assert snapshots() == wanted
|
||
|
|
for row in wanted["research_revisions"]:
|
||
|
|
(TemplateSpec if row["asset_id"] == "a" else FeatureSpec).model_validate(row["content"])
|
||
|
|
# Downgrade does not invent the deleted metadata; re-upgrade is idempotent.
|
||
|
|
command.downgrade(config, "0022")
|
||
|
|
assert snapshots() == wanted
|
||
|
|
command.upgrade(config, "head")
|
||
|
|
assert snapshots() == wanted
|
||
|
|
command.check(config)
|
||
|
|
engine.dispose()
|