mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-13 06:31:47 +00:00
9699f04cdc
* feat(projects): add visible project memory with extraction controls * fix(projects): avoid replaying committed memory extractions * chore(ci): retrigger project memory checks --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
70 lines
3.9 KiB
Python
70 lines
3.9 KiB
Python
"""Both deployed migration branches converge without losing Run data or policy."""
|
|
|
|
import asyncio
|
|
import os
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from alembic import command
|
|
from sqlalchemy import inspect, text
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
from support.postgres import asyncpg_test_url
|
|
|
|
from deerflow.persistence import bootstrap
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("backend", ["sqlite", "postgres"])
|
|
@pytest.mark.parametrize("previous", ["0049_run_inputs", "0049_run_producer_fencing", "0050_history_seed_authority"])
|
|
async def test_both_migration_branches_upgrade_to_one_head(tmp_path, backend, previous):
|
|
schema = None
|
|
if backend == "postgres":
|
|
uri = os.environ.get("TEST_POSTGRES_URI")
|
|
if not uri:
|
|
pytest.skip("requires TEST_POSTGRES_URI")
|
|
schema = "merged_runs_" + uuid4().hex
|
|
engine = create_async_engine(asyncpg_test_url(uri), connect_args={"server_settings": {"search_path": schema}})
|
|
else:
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}")
|
|
cfg = bootstrap._get_alembic_config(engine, postgres_schema=schema or "")
|
|
try:
|
|
if schema:
|
|
async with engine.begin() as conn:
|
|
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
|
await asyncio.to_thread(command.upgrade, cfg, previous)
|
|
async with engine.begin() as conn:
|
|
columns = await conn.run_sync(lambda sync: {c["name"] for c in inspect(sync).get_columns("runs")})
|
|
if previous == "0049_run_inputs":
|
|
assert "run_inputs_json" in columns
|
|
assert "run_incarnation" not in columns
|
|
assert not await conn.run_sync(lambda sync: inspect(sync).has_table("run_producer_policy"))
|
|
else:
|
|
assert "run_inputs_json" not in columns
|
|
assert {"run_incarnation", "producer_epoch", "producer_state"} <= columns
|
|
if schema:
|
|
await conn.execute(text("INSERT INTO run_producer_policy(id,strict,cutover_id) VALUES (1,true,'preserved-policy')"))
|
|
await conn.execute(
|
|
text("""INSERT INTO runs(
|
|
run_id,thread_id,user_id,status,multitask_strategy,metadata_json,kwargs_json,
|
|
message_count,total_input_tokens,total_output_tokens,total_tokens,llm_call_count,
|
|
lead_agent_tokens,subagent_tokens,middleware_tokens,created_at,updated_at
|
|
) VALUES ('old','thread','alice','success','reject','{}','{}',0,0,0,0,0,0,0,0,CURRENT_TIMESTAMP,CURRENT_TIMESTAMP)""")
|
|
)
|
|
await asyncio.to_thread(command.upgrade, cfg, "head")
|
|
async with engine.connect() as conn:
|
|
assert (await conn.execute(text("SELECT version_num FROM alembic_version"))).scalars().all() == ["0054_project_memory"]
|
|
columns = await conn.run_sync(lambda sync: {c["name"] for c in inspect(sync).get_columns("runs")})
|
|
assert {"run_inputs_json", "run_incarnation", "producer_epoch", "producer_state"} <= columns
|
|
assert (await conn.execute(text("SELECT status,run_inputs_json,run_incarnation,producer_epoch,producer_state FROM runs WHERE run_id='old'"))).one() == ("success", None, None, None, None)
|
|
if schema:
|
|
assert await conn.scalar(text("SELECT count(*) FROM pg_trigger WHERE tgrelid='run_events'::regclass AND tgname='run_event_producer_guard'")) == 1
|
|
if previous != "0049_run_inputs":
|
|
assert (await conn.execute(text("SELECT strict,cutover_id FROM run_producer_policy WHERE id=1"))).one() == (True, "preserved-policy")
|
|
# Bootstrap must also recognize the merged stamp on a later restart.
|
|
await bootstrap.bootstrap_schema(engine, backend=backend, postgres_schema=schema or "")
|
|
finally:
|
|
if schema:
|
|
async with engine.begin() as conn:
|
|
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
|
await engine.dispose()
|