92 lines
3.6 KiB
Python
92 lines
3.6 KiB
Python
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from deerflow.config.checkpointer_config import CheckpointerConfig
|
|
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
|
|
from deerflow.runtime.store.async_provider import make_store
|
|
from deerflow.runtime.runs.manager import RunRecord
|
|
from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus
|
|
from deerflow.runtime.runs.worker import _finalize_run
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_sqlite_initialization_and_shared_store(tmp_path):
|
|
config = SimpleNamespace(checkpointer=CheckpointerConfig(type="sqlite", connection_string=str(tmp_path / "chat.db")))
|
|
|
|
async def node(index):
|
|
async with make_checkpointer(config) as saver, make_store(config) as store:
|
|
assert (await (await saver.conn.execute("PRAGMA busy_timeout")).fetchone())[0] == 30000
|
|
await store.aput(("qa",), str(index), {"index": index})
|
|
assert (await store.aget(("qa",), str(index))).value == {"index": index}
|
|
|
|
await asyncio.gather(*(node(i) for i in range(4)))
|
|
async with make_store(config) as store:
|
|
assert len(await store.asearch(("qa",))) == 4
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pg_disabled_checkpoint_and_store_use_same_sqlite(tmp_path, monkeypatch):
|
|
monkeypatch.setenv("DEER_FLOW_USE_PG_CHECKPOINTER", "false")
|
|
path = str(tmp_path / "fallback.db")
|
|
config = SimpleNamespace(checkpointer=CheckpointerConfig(type="postgres", connection_string="postgresql://unused", sqlite_connection_string=path))
|
|
async with make_checkpointer(config) as saver, make_store(config) as store:
|
|
saver_path = (await (await saver.conn.execute("PRAGMA database_list")).fetchone())[2]
|
|
store_path = (await (await store.conn.execute("PRAGMA database_list")).fetchone())[2]
|
|
assert saver_path == store_path
|
|
await store.aput(("qa",), "fallback", {"ok": True})
|
|
assert (await store.aget(("qa",), "fallback")).value == {"ok": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_finalization_does_not_wait_for_title_checkpoint_read():
|
|
title_read_started = asyncio.Event()
|
|
release_title_read = asyncio.Event()
|
|
status_updates: list[tuple[str, str]] = []
|
|
titles: list[tuple[str, str]] = []
|
|
|
|
class Checkpointer:
|
|
async def aget_tuple(self, _config):
|
|
title_read_started.set()
|
|
await release_title_read.wait()
|
|
return SimpleNamespace(checkpoint={"channel_values": {"title": "异步标题"}})
|
|
|
|
class ThreadStore:
|
|
async def update_status(self, thread_id: str, status: str) -> None:
|
|
status_updates.append((thread_id, status))
|
|
|
|
async def update_display_name(self, thread_id: str, title: str) -> None:
|
|
titles.append((thread_id, title))
|
|
|
|
class Manager:
|
|
def is_latest(self, _record: RunRecord) -> bool:
|
|
return True
|
|
|
|
record = RunRecord(
|
|
run_id="run-title-test",
|
|
thread_id="thread-title-test",
|
|
assistant_id="assistant",
|
|
status=RunStatus.success,
|
|
on_disconnect=DisconnectMode.continue_,
|
|
)
|
|
await asyncio.wait_for(
|
|
_finalize_run(
|
|
journal=None,
|
|
run_manager=Manager(),
|
|
record=record,
|
|
ctx=SimpleNamespace(checkpointer=Checkpointer(), thread_store=ThreadStore(), app_config=None),
|
|
),
|
|
timeout=0.1,
|
|
)
|
|
|
|
assert status_updates == [("thread-title-test", "idle")]
|
|
await asyncio.wait_for(title_read_started.wait(), timeout=0.1)
|
|
assert titles == []
|
|
release_title_read.set()
|
|
for _ in range(10):
|
|
if titles:
|
|
break
|
|
await asyncio.sleep(0)
|
|
assert titles == [("thread-title-test", "异步标题")]
|