269 lines
9.0 KiB
Python
269 lines
9.0 KiB
Python
"""Shared behavior tests for the artifact-library read index."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.persistence.artifact_library import (
|
|
ArtifactLibraryItemRow,
|
|
ArtifactLibraryStore,
|
|
ArtifactLibrarySyncStateRow,
|
|
MemoryArtifactLibraryStore,
|
|
SqlArtifactLibraryStore,
|
|
)
|
|
|
|
|
|
@pytest.fixture(params=("memory", "sqlite"), ids=("memory", "sqlite"))
|
|
def store_backend(request: pytest.FixtureRequest) -> str:
|
|
return str(request.param)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _store_for(
|
|
backend: str,
|
|
tmp_path: Path,
|
|
) -> AsyncIterator[ArtifactLibraryStore]:
|
|
if backend == "memory":
|
|
yield MemoryArtifactLibraryStore()
|
|
return
|
|
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'artifact-library.db'}")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(ArtifactLibrarySyncStateRow.__table__.create)
|
|
await connection.run_sync(ArtifactLibraryItemRow.__table__.create)
|
|
try:
|
|
yield SqlArtifactLibraryStore(async_sessionmaker(engine, expire_on_commit=False))
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
def _item(
|
|
index: int,
|
|
name: str,
|
|
*,
|
|
thread_id: str | None = None,
|
|
path: str | None = None,
|
|
file_kind: str = "document",
|
|
source_type: str = "general_qa",
|
|
thread_title: str | None = None,
|
|
agent_name: str | None = None,
|
|
strategy_name: str | None = None,
|
|
strategy_project_id: str | None = None,
|
|
) -> dict:
|
|
return {
|
|
"thread_id": thread_id or f"thread-{index}",
|
|
"thread_title": thread_title or f"Conversation {index}",
|
|
"agent_id": f"agent-{index}",
|
|
"agent_name": agent_name,
|
|
"thread_type": "agent" if agent_name else "general",
|
|
"source_type": source_type,
|
|
"strategy_project_id": strategy_project_id,
|
|
"strategy_name": strategy_name,
|
|
"strategy_stale": False,
|
|
"strategy_source_kind": "synthesis" if strategy_name else None,
|
|
"download_url": f"/api/artifacts/{index}/{name}",
|
|
"path": path or f"outputs/{name}",
|
|
"name": name,
|
|
"size_bytes": index * 100,
|
|
"mime_type": "text/plain",
|
|
"modified_at": datetime(2026, 1, 1, 0, index, tzinfo=UTC).isoformat(),
|
|
"file_kind": file_kind,
|
|
}
|
|
|
|
|
|
async def _publish(
|
|
store: ArtifactLibraryStore,
|
|
user_id: str,
|
|
items: list[dict],
|
|
) -> str:
|
|
generation = await store.claim_sync(user_id, force=True)
|
|
assert generation is not None
|
|
assert await store.finish_sync(user_id, generation, items) is True
|
|
return generation
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_second_page_uses_real_offset(
|
|
store_backend: str,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
async with _store_for(store_backend, tmp_path) as store:
|
|
await _publish(
|
|
store,
|
|
"user-a",
|
|
[_item(index, f"file-{index}.txt") for index in range(5)],
|
|
)
|
|
|
|
first_total, first_page = await store.list_items("user-a", offset=0, limit=2)
|
|
second_total, second_page = await store.list_items("user-a", offset=2, limit=2)
|
|
|
|
assert first_total == second_total == 5
|
|
assert [row["name"] for row in first_page] == [
|
|
"file-4.txt",
|
|
"file-3.txt",
|
|
]
|
|
assert [row["name"] for row in second_page] == [
|
|
"file-2.txt",
|
|
"file-1.txt",
|
|
]
|
|
assert {row["path"] for row in first_page}.isdisjoint(row["path"] for row in second_page)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filters_search_and_user_isolation(
|
|
store_backend: str,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
async with _store_for(store_backend, tmp_path) as store:
|
|
await _publish(
|
|
store,
|
|
"user-a",
|
|
[
|
|
_item(1, "brief.txt", agent_name="Orbit Analyst"),
|
|
_item(
|
|
2,
|
|
"chart.png",
|
|
file_kind="image",
|
|
source_type="strategy_selection",
|
|
strategy_name="North Wind",
|
|
strategy_project_id="strategy-42",
|
|
),
|
|
_item(3, "bundle.zip", file_kind="archive"),
|
|
],
|
|
)
|
|
await _publish(
|
|
store,
|
|
"user-b",
|
|
[
|
|
_item(
|
|
9,
|
|
"private.txt",
|
|
thread_id="thread-1",
|
|
path="outputs/brief.txt",
|
|
agent_name="Orbit Analyst",
|
|
)
|
|
],
|
|
)
|
|
|
|
total, rows = await store.list_items(
|
|
"user-a",
|
|
query="NORTH WIND",
|
|
file_kind="image",
|
|
source_type="strategy_selection",
|
|
)
|
|
assert total == 1
|
|
assert [row["name"] for row in rows] == ["chart.png"]
|
|
|
|
total, rows = await store.list_items("user-a", query="orbit analyst")
|
|
assert total == 1
|
|
assert [row["name"] for row in rows] == ["brief.txt"]
|
|
|
|
total, rows = await store.list_items("user-a")
|
|
assert total == 3
|
|
assert "private.txt" not in {row["name"] for row in rows}
|
|
total, rows = await store.list_items("user-b")
|
|
assert total == 1
|
|
assert [row["name"] for row in rows] == ["private.txt"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_new_generation_replaces_rows_and_deletes_stale_items(
|
|
store_backend: str,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
async with _store_for(store_backend, tmp_path) as store:
|
|
retained = _item(1, "retained-v1.txt", path="outputs/retained.txt")
|
|
stale = _item(2, "stale.txt")
|
|
first_generation = await _publish(
|
|
store,
|
|
"user-a",
|
|
[retained, stale],
|
|
)
|
|
|
|
second_generation = await store.claim_sync("user-a", force=True)
|
|
assert second_generation is not None
|
|
assert second_generation != first_generation
|
|
# Duplicate identities in one snapshot are deterministic: the last
|
|
# observation wins, while identities absent from this generation go.
|
|
retained_v2 = {
|
|
**retained,
|
|
"name": "retained-v2.txt",
|
|
"modified_at": datetime(2026, 1, 1, 1, 0, tzinfo=UTC).isoformat(),
|
|
}
|
|
retained_v3 = {**retained_v2, "name": "retained-v3.txt"}
|
|
assert await store.finish_sync(
|
|
"user-a",
|
|
second_generation,
|
|
[retained_v2, retained_v3, _item(3, "new.txt")],
|
|
)
|
|
|
|
total, rows = await store.list_items("user-a")
|
|
assert total == 2
|
|
assert {row["name"] for row in rows} == {
|
|
"retained-v3.txt",
|
|
"new.txt",
|
|
}
|
|
assert "stale.txt" not in {row["name"] for row in rows}
|
|
|
|
failed_generation = await store.claim_sync("user-a", force=True)
|
|
assert failed_generation is not None
|
|
assert await store.fail_sync("user-a", failed_generation, "filesystem unavailable")
|
|
# A failed reconciliation changes only its state; the last complete
|
|
# generation remains queryable instead of publishing a partial page.
|
|
total, rows = await store.list_items("user-a")
|
|
assert total == 2
|
|
assert {row["name"] for row in rows} == {
|
|
"retained-v3.txt",
|
|
"new.txt",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_live_lease_is_deduplicated_and_generation_is_fenced(
|
|
store_backend: str,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
async with _store_for(store_backend, tmp_path) as store:
|
|
claims = await asyncio.gather(
|
|
*(
|
|
store.claim_sync(
|
|
"user-a",
|
|
force=True,
|
|
lease_seconds=60,
|
|
)
|
|
for _ in range(4)
|
|
)
|
|
)
|
|
winners = [generation for generation in claims if generation is not None]
|
|
assert len(winners) == 1
|
|
winner = winners[0]
|
|
assert await store.claim_sync("user-a", force=True) is None
|
|
state = await store.get_sync_state("user-a")
|
|
assert state["status"] == "syncing"
|
|
assert state["generation"] == winner
|
|
|
|
assert await store.finish_sync("user-a", winner, [_item(1, "winner.txt")])
|
|
assert await store.finish_sync("user-a", winner, [_item(2, "late.txt")]) is False
|
|
assert await store.fail_sync("user-a", winner, "late failure") is False
|
|
|
|
expired = await store.claim_sync(
|
|
"user-a",
|
|
force=True,
|
|
lease_seconds=0,
|
|
)
|
|
assert expired is not None
|
|
replacement = await store.claim_sync("user-a", force=True)
|
|
assert replacement is not None and replacement != expired
|
|
assert await store.finish_sync("user-a", expired, [_item(3, "stale.txt")]) is False
|
|
assert await store.finish_sync("user-a", replacement, [_item(4, "replacement.txt")])
|
|
total, rows = await store.list_items("user-a")
|
|
assert total == 1
|
|
assert [row["name"] for row in rows] == ["replacement.txt"]
|