import asyncio import importlib import os import sys import threading from pathlib import Path from fastapi.testclient import TestClient os.environ["JOB_BACKEND"] = "memory" os.environ["DOCS_BACKEND"] = "memory" CURRENT_DIR = Path(__file__).resolve().parent BACKEND_DIR = CURRENT_DIR.parent if str(BACKEND_DIR) not in sys.path: sys.path.insert(0, str(BACKEND_DIR)) import job_handlers # type: ignore import job_system # type: ignore import risk_control # type: ignore import session_store # type: ignore import audit_store # type: ignore main = importlib.import_module("main") HEADERS = {"X-API-Key": main.API_KEY} def setup_function(): job_system.reset_job_manager() risk_control.reset_risk_controller() session_store.reset_session_store() audit_store.reset_audit_store() main._handlers_registered = False def _payload(): return { "prefix": "比较当前主流向量数据库的设计差异", "suffix": "", "languageId": "markdown", "privacy_mode": True, } def test_is_blocked_public_url(): assert job_handlers._is_blocked_public_url("http://127.0.0.1/test") is True assert job_handlers._is_blocked_public_url("file:///tmp/test") is True assert job_handlers._is_blocked_public_url("https://example.com/docs") is False def test_web_search_route_returns_done(monkeypatch): async def fake_call_ollama(prompt, system_prompt=None, tag="", **kwargs): # noqa: ARG001 if tag.endswith("-webq"): return {"content": '["vector database comparison", "pinecone weaviate qdrant"]'} if tag.endswith("-webu"): return {"content": '["https://example.com/a", "https://example.com/b"]'} if tag.endswith("-webf"): return {"content": "第一段\n\n第二段"} raise AssertionError(f"unexpected tag: {tag}") async def fake_searxng_search(query, *, limit): # noqa: ARG001 return [ { "title": "Doc A", "url": "https://example.com/a", "score": 9.1, "published_date": "2026-06-08", "snippet": "snippet a", }, { "title": "Doc B", "url": "https://example.com/b", "score": 8.8, "published_date": "2026-06-07", "snippet": "snippet b", }, ] async def fake_firecrawl_scrape(url): return {"url": url, "title": f"title for {url}", "markdown": f"content for {url}"} monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama) monkeypatch.setattr(job_handlers, "_searxng_search", fake_searxng_search) monkeypatch.setattr(job_handlers, "_firecrawl_scrape", fake_firecrawl_scrape) with TestClient(main.app) as client: with client.stream("POST", "/v1/web-search", headers=HEADERS, json=_payload()) as resp: assert resp.status_code == 200 body = "".join(resp.iter_text()) assert "event: progress" in body assert "keywords" in body assert "searching" in body assert "selecting_urls" in body assert "crawling" in body assert "synthesizing" in body assert "event: done" in body assert "第一段" in body def test_cancel_web_search(monkeypatch): started = threading.Event() cancelled = threading.Event() async def fake_call_ollama(*args, **kwargs): tag = kwargs.get("tag", "") if tag.endswith("-webq"): started.set() try: while True: await asyncio.sleep(0.05) except asyncio.CancelledError: cancelled.set() raise return {"content": "[]"} monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama) request_id = "req-web-search-cancel" with TestClient(main.app) as client: response_box = {} def send_request(): with client.stream( "POST", "/v1/web-search", headers={**HEADERS, "X-Request-Id": request_id}, json=_payload(), ) as response: response_box["status_code"] = response.status_code response_box["body"] = "".join(response.iter_text()) search_thread = threading.Thread(target=send_request, daemon=True) search_thread.start() assert started.wait(timeout=2.0) cancel_response = client.post( "/v1/web-search/cancel", headers=HEADERS, json={"request_id": request_id, "reason": "abort"}, ) assert cancel_response.status_code == 200 assert cancel_response.json() == {"cancelled": True, "status": "ok"} search_thread.join(timeout=5.0) assert not search_thread.is_alive() assert cancelled.wait(timeout=2.0) assert "event: cancelled" in response_box["body"]