mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-31 15:40:37 +00:00
171 lines
8.5 KiB
Python
171 lines
8.5 KiB
Python
"""Bounded, request-local detail checks retain WeKnora's authorization gate."""
|
|
|
|
import asyncio
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from deerflow.knowledge import KnowledgeRequestContext
|
|
|
|
|
|
@pytest.fixture
|
|
def setup_provider(monkeypatch):
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "examples/deerflow-extension-weknora"))
|
|
from deerflow_extension_weknora import WeKnoraProvider
|
|
|
|
monkeypatch.setenv("WK_CONCURRENCY_KEY", "synthetic")
|
|
|
|
def create(*, gate=False, failure=None, hold=False):
|
|
state = {"active": 0, "peak": 0, "calls": Counter(), "searches": 0}
|
|
started, release = asyncio.Event(), asyncio.Event()
|
|
|
|
async def handle(request):
|
|
assert request.headers["X-API-Key"] == "synthetic"
|
|
assert request.headers["X-Tenant-ID"] == "1"
|
|
path = request.url.path
|
|
if path == "/api/v1/knowledge-bases":
|
|
data = [{"id": "kb", "name": "Synthetic"}]
|
|
elif path.endswith("hybrid-search"):
|
|
state["searches"] += 1
|
|
data = [{"knowledge_id": f"doc-{i}", "content": f"chunk {i}"} for i in range(8)]
|
|
data.insert(1, {"knowledge_id": "doc-0", "content": "second chunk"})
|
|
else:
|
|
doc = path.rsplit("/", 1)[1]
|
|
state["calls"][doc] += 1
|
|
state["active"] += 1
|
|
state["peak"] = max(state["peak"], state["active"])
|
|
started.set()
|
|
try:
|
|
if state["active"] == 4:
|
|
release.set()
|
|
if gate:
|
|
await asyncio.wait_for(release.wait(), 2)
|
|
if failure and doc == "doc-0":
|
|
return httpx.Response(failure, json={"success": False})
|
|
if hold:
|
|
await asyncio.Event().wait()
|
|
await asyncio.sleep(0)
|
|
data = {"id": doc, "knowledge_base_id": "kb", "parse_status": "completed", "enable_status": "enabled", "title": doc}
|
|
finally:
|
|
state["active"] -= 1
|
|
return httpx.Response(200, json={"success": True, "data": data})
|
|
|
|
provider = WeKnoraProvider({"spaces": [{"base_url": "https://weknora.invalid", "tenant_id": "1", "api_key_env": "WK_CONCURRENCY_KEY", "user_ids": ["alice"]}]}, transport=httpx.MockTransport(handle))
|
|
return provider, state, started
|
|
|
|
return create
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("selected", [False, True])
|
|
async def test_details_overlap_with_a_bound_and_preserve_hit_order(setup_provider, selected):
|
|
provider, state, _ = setup_provider(gate=True)
|
|
ctx = KnowledgeRequestContext("alice", "lead_agent")
|
|
if selected:
|
|
catalog = await provider.catalog(context=ctx, page=1, page_size=20, search="")
|
|
dataset = catalog.items[0].id
|
|
ctx = KnowledgeRequestContext("alice", "lead_agent", {"version": 1, "mode": "selected", "dataset_ids": [dataset], "document_filters": [{"dataset_id": dataset, "document_ids": [f"doc-{i}" for i in range(8)]}]})
|
|
result = await provider.search_scoped("q", context=ctx)
|
|
assert state["peak"] == 4
|
|
assert state["active"] == 0
|
|
assert state["calls"] == Counter({f"doc-{i}": 1 for i in range(8)})
|
|
assert [h.content for h in result.hits] == ["chunk 0", "second chunk", *[f"chunk {i}" for i in range(1, 8)]]
|
|
assert [h.source.title for h in result.hits] == ["doc-0", "doc-0", *[f"doc-{i}" for i in range(1, 8)]]
|
|
# A second call must revalidate documents rather than reuse authorization.
|
|
await provider.search_scoped("q", context=ctx)
|
|
assert all(n == 2 for n in state["calls"].values())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure,error", [(403, PermissionError), (404, FileNotFoundError), (500, RuntimeError)])
|
|
async def test_failure_drains_other_detail_requests_before_returning(setup_provider, failure, error):
|
|
provider, state, _ = setup_provider(gate=True, failure=failure, hold=True)
|
|
with pytest.raises(error):
|
|
await asyncio.wait_for(provider.search_scoped("q", context=KnowledgeRequestContext("alice")), 5)
|
|
assert state["peak"] == 4
|
|
assert state["active"] == 0
|
|
assert len(state["calls"]) == 4 # Later batches never start after failure.
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancellation_drains_detail_requests(setup_provider):
|
|
provider, state, started = setup_provider(hold=True)
|
|
task = asyncio.create_task(provider.search_scoped("q", context=KnowledgeRequestContext("alice")))
|
|
await asyncio.wait_for(started.wait(), 5)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert state["active"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_selected_document_denial_prevents_search(setup_provider):
|
|
provider, state, _ = setup_provider(gate=True, failure=403, hold=True)
|
|
catalog = await provider.catalog(context=KnowledgeRequestContext("alice"), page=1, page_size=20, search="")
|
|
dataset = catalog.items[0].id
|
|
ctx = KnowledgeRequestContext("alice", "lead_agent", {"version": 1, "mode": "selected", "dataset_ids": [dataset], "document_filters": [{"dataset_id": dataset, "document_ids": [f"doc-{i}" for i in range(8)]}]})
|
|
with pytest.raises(PermissionError):
|
|
await asyncio.wait_for(provider.search_scoped("q", context=ctx), 5)
|
|
assert state["searches"] == 0
|
|
assert state["active"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_equal_native_ids_in_different_tenants_are_revalidated(monkeypatch):
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "examples/deerflow-extension-weknora"))
|
|
from deerflow_extension_weknora import WeKnoraProvider
|
|
|
|
monkeypatch.setenv("WK_ONE", "key-one")
|
|
monkeypatch.setenv("WK_TWO", "key-two")
|
|
details = []
|
|
|
|
async def handle(request):
|
|
tenant = request.headers["X-Tenant-ID"]
|
|
assert request.headers["X-API-Key"] == {"1": "key-one", "2": "key-two"}[tenant]
|
|
if request.url.path == "/api/v1/knowledge-bases":
|
|
data = [{"id": "kb", "name": tenant}]
|
|
elif request.url.path.endswith("hybrid-search"):
|
|
data = [{"knowledge_id": "same-id", "content": f"evidence {tenant}"}]
|
|
else:
|
|
details.append(tenant)
|
|
data = {"id": "same-id", "knowledge_base_id": "kb", "parse_status": "completed", "enable_status": "enabled", "title": tenant}
|
|
return httpx.Response(200, json={"success": True, "data": data})
|
|
|
|
provider = WeKnoraProvider(
|
|
{"spaces": [{"base_url": "https://weknora.invalid", "tenant_id": tenant, "api_key_env": key, "user_ids": ["alice"]} for tenant, key in [("1", "WK_ONE"), ("2", "WK_TWO")]]}, transport=httpx.MockTransport(handle)
|
|
)
|
|
result = await provider.search_scoped("q", context=KnowledgeRequestContext("alice"))
|
|
assert details == ["1", "2"]
|
|
assert [hit.source.title for hit in result.hits] == ["1", "2"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_result_cap_checks_empty_hits_but_not_the_unconsumed_suffix(monkeypatch):
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "examples/deerflow-extension-weknora"))
|
|
from deerflow_extension_weknora import WeKnoraProvider
|
|
|
|
monkeypatch.setenv("WK_CAP_KEY", "synthetic")
|
|
details = []
|
|
|
|
async def handle(request):
|
|
path = request.url.path
|
|
if path == "/api/v1/knowledge-bases":
|
|
data = [{"id": "first", "name": "First"}, {"id": "last", "name": "Last"}]
|
|
elif path.endswith("first/hybrid-search"):
|
|
data = [{"knowledge_id": "first-doc", "content": f"chunk {i}"} for i in range(99)]
|
|
elif path.endswith("last/hybrid-search"):
|
|
data = [{"knowledge_id": "empty-doc", "content": ""}, {"knowledge_id": "last-doc", "content": "last chunk"}, {"knowledge_id": "unsafe/path", "content": "unconsumed"}]
|
|
else:
|
|
doc = path.rsplit("/", 1)[1]
|
|
details.append(doc)
|
|
data = {"id": doc, "knowledge_base_id": "first" if doc == "first-doc" else "last", "parse_status": "completed", "enable_status": "enabled", "title": doc}
|
|
return httpx.Response(200, json={"success": True, "data": data})
|
|
|
|
provider = WeKnoraProvider({"spaces": [{"base_url": "https://weknora.invalid", "tenant_id": "1", "api_key_env": "WK_CAP_KEY", "user_ids": ["alice"]}]}, transport=httpx.MockTransport(handle))
|
|
result = await provider.search_scoped("q", context=KnowledgeRequestContext("alice"))
|
|
assert len(result.hits) == 100
|
|
assert result.hits[-1].content == "last chunk"
|
|
assert details == ["first-doc", "empty-doc", "last-doc"]
|