152 lines
6.4 KiB
Python
152 lines
6.4 KiB
Python
"""Regression checks for the ordinary knowledge-base initial request fan-out."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException
|
|
|
|
from deerflow.integrations.weknora.client import WeKnoraClient, WeKnoraError
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unavailable_directory_fails_concurrent_callers_once() -> None:
|
|
calls = 0
|
|
|
|
async def upstream(_: httpx.Request) -> httpx.Response:
|
|
nonlocal calls
|
|
calls += 1
|
|
await asyncio.sleep(0)
|
|
raise httpx.ReadTimeout("unavailable")
|
|
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(upstream)) as transport:
|
|
client = WeKnoraClient(base_url="http://weknora.test", http_client=transport)
|
|
results = await asyncio.gather(*(client.list_knowledge_bases() for _ in range(6)), return_exceptions=True)
|
|
assert calls == 1
|
|
assert all(isinstance(result, WeKnoraError) and result.status_code == 504 for result in results)
|
|
with pytest.raises(WeKnoraError):
|
|
await client.list_knowledge_bases()
|
|
assert calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_write_directory_does_not_join_or_cache_an_older_inflight_read() -> None:
|
|
started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
calls = 0
|
|
|
|
async def upstream(request: httpx.Request) -> httpx.Response:
|
|
nonlocal calls
|
|
if request.method == "PUT":
|
|
return httpx.Response(200, json={"data": {"id": "base", "name": "New"}})
|
|
calls += 1
|
|
if calls == 1:
|
|
started.set()
|
|
await release.wait()
|
|
return httpx.Response(200, json={"data": [{"id": "base", "name": "Old"}]})
|
|
return httpx.Response(200, json={"data": [{"id": "base", "name": "New"}]})
|
|
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(upstream)) as transport:
|
|
client = WeKnoraClient(base_url="http://weknora.test", http_client=transport)
|
|
old_read = asyncio.create_task(client.list_knowledge_bases())
|
|
await started.wait()
|
|
await client.update_knowledge_base("base", {"name": "New"})
|
|
try:
|
|
new = await asyncio.wait_for(client.list_knowledge_bases(), timeout=1)
|
|
finally:
|
|
release.set()
|
|
old = await old_read
|
|
cached = await client.list_knowledge_bases()
|
|
assert old[0]["name"] == "Old"
|
|
assert new == cached == [{"id": "base", "name": "New"}]
|
|
assert calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("status", [200, 503])
|
|
async def test_iframe_startup_coalesces_admin_logins_and_failures(monkeypatch: pytest.MonkeyPatch, status: int) -> None:
|
|
from app.gateway.routers import llmwiki
|
|
|
|
runtime = SimpleNamespace(api_base_url="http://weknora-login.test", admin_email="admin@test", admin_password="test-only")
|
|
calls = 0
|
|
|
|
async def upstream(_: httpx.Request) -> httpx.Response:
|
|
nonlocal calls
|
|
calls += 1
|
|
await asyncio.sleep(0)
|
|
return httpx.Response(status, json={"data": {"token": "token", "tenant": {"id": "tenant"}}})
|
|
|
|
real_client = httpx.AsyncClient
|
|
monkeypatch.setattr(llmwiki, "_WEKNORA_ADMIN_SESSION_CACHE", {})
|
|
monkeypatch.setattr(llmwiki, "_WEKNORA_ADMIN_SESSION_TASKS", {})
|
|
monkeypatch.setattr(llmwiki.httpx, "AsyncClient", lambda **kwargs: real_client(transport=httpx.MockTransport(upstream), **kwargs))
|
|
results = await asyncio.gather(*(llmwiki._get_weknora_admin_session(runtime) for _ in range(6)), return_exceptions=True)
|
|
|
|
assert calls == 1
|
|
if status == 200:
|
|
assert all(result["token"] == "token" for result in results)
|
|
assert (await llmwiki._get_weknora_admin_session(runtime))["token"] == "token"
|
|
assert calls == 1
|
|
else:
|
|
assert all(isinstance(result, HTTPException) and result.status_code == 502 for result in results)
|
|
# Failed logins are not cached, so a later navigation can recover.
|
|
with pytest.raises(HTTPException):
|
|
await llmwiki._get_weknora_admin_session(runtime)
|
|
assert calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_iframe_proxy_shares_one_connection_pool_and_closes_it_on_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from app.gateway.routers import llmwiki
|
|
|
|
runtime = SimpleNamespace(web_base_url="http://weknora-web.test", api_base_url="http://weknora-api.test")
|
|
ctx = llmwiki.WeKnoraEmbedContext(mapping_id="local", weknora_id="remote", user_id="owner", is_admin=False, can_write=True, exp=4_102_444_800)
|
|
pools: list[httpx.AsyncClient] = []
|
|
requests: list[str] = []
|
|
real_client = httpx.AsyncClient
|
|
|
|
async def upstream(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request.url.path)
|
|
assert "cookie" not in request.headers
|
|
await asyncio.sleep(0)
|
|
return httpx.Response(
|
|
200,
|
|
text="export default {};",
|
|
headers={"content-type": "application/javascript", "set-cookie": "upstream-session=must-not-be-shared; Path=/"},
|
|
)
|
|
|
|
def create_pool(**kwargs) -> httpx.AsyncClient:
|
|
client = real_client(transport=httpx.MockTransport(upstream), **kwargs)
|
|
pools.append(client)
|
|
return client
|
|
|
|
async def context(_):
|
|
return ctx
|
|
|
|
async def session(_, **__):
|
|
return {"token": "token", "tenant": {}}
|
|
|
|
monkeypatch.setattr(llmwiki.httpx, "AsyncClient", create_pool)
|
|
monkeypatch.setattr(llmwiki, "_runtime_for_iframe_proxy", lambda _: runtime)
|
|
monkeypatch.setattr(llmwiki, "_read_embed_context", context)
|
|
monkeypatch.setattr(llmwiki, "_verify_embed_cookie", lambda _: ctx)
|
|
monkeypatch.setattr(llmwiki, "_get_weknora_admin_session", session)
|
|
app = FastAPI()
|
|
app.include_router(llmwiki.router)
|
|
app.include_router(llmwiki.proxy_router)
|
|
|
|
async with app.router.lifespan_context(app):
|
|
async with real_client(transport=httpx.ASGITransport(app=app), base_url="http://gateway.test") as client:
|
|
responses = await asyncio.gather(client.get("/assets/one.js"), client.get("/assets/two.js"))
|
|
responses.append(await client.get("/assets/three.js", headers={"Cookie": "browser-session=must-not-be-forwarded"}))
|
|
assert all(response.status_code == 200 for response in responses)
|
|
assert len(pools) == 1
|
|
assert not pools[0].is_closed
|
|
assert not pools[0].cookies
|
|
assert sorted(requests) == ["/assets/one.js", "/assets/three.js", "/assets/two.js"]
|
|
assert pools[0].is_closed
|
|
assert app.state.llmwiki_proxy_http_client is None
|