llm-in-text/backend/tests/test_web_search.py

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"]