deer-flow/backend/tests/test_knowledge_router.py

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest
from _router_auth_helpers import make_authed_test_app
from fastapi.testclient import TestClient

from app.gateway.auth.models import User
from app.gateway.deps import get_config
from app.gateway.routers import knowledge


def _config(
    *,
    enabled: bool = True,
    api_key: str | None = "ragflow-secret",
    scope_selection_enabled: bool = False,
    datasets: list[str] | None = None,
    provider: str = "deerflow.community.ragflow.tools:knowledge_search_tool",
) -> SimpleNamespace:
    tool = SimpleNamespace(
        use=provider,
        model_extra={
            "base_url": "http://ragflow.test",
            "api_key": api_key,
            "timeout": 30,
            **({"datasets": datasets} if datasets is not None else {}),
        },
    )
    return SimpleNamespace(
        knowledge_base=SimpleNamespace(
            enabled=enabled,
            scope_selection_enabled=scope_selection_enabled,
        ),
        get_tool_config=lambda name: tool if name == "knowledge_search" else None,
    )


def _user() -> User:
    return User(
        email="router-test@example.com",
        password_hash="x",
        system_role="user",
    )


def _app(
    monkeypatch: pytest.MonkeyPatch,
    client: object,
    *,
    config: SimpleNamespace | None = None,
):
    app = make_authed_test_app(user_factory=_user)
    app.include_router(knowledge.router)
    app.dependency_overrides[get_config] = lambda: config or _config()
    monkeypatch.setattr(
        knowledge,
        "_build_retrieval_client",
        lambda settings: client,
    )
    return app


def _enable_scope_catalog(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(
        knowledge,
        "load_agent_config",
        lambda name, *, user_id: SimpleNamespace(
            name=name,
            tool_groups=["knowledge"],
        ),
    )


def test_retrieval_catalog_enforces_allowlist_and_normalizes_pages(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _enable_scope_catalog(monkeypatch)
    ragflow = SimpleNamespace(
        list_datasets=AsyncMock(
            side_effect=[
                [
                    {
                        "id": "dataset-1",
                        "name": "Policies",
                        "embedding_model": "embed-a",
                        "chunk_count": 3,
                    }
                ],
                [
                    {
                        "id": "dataset-2",
                        "name": "Empty",
                        "embedding_model": "",
                        "chunk_count": 0,
                    }
                ],
            ]
        )
    )
    config = _config(
        scope_selection_enabled=True,
        datasets=["dataset-1", "dataset-2"],
    )

    with TestClient(_app(monkeypatch, ragflow, config=config)) as client:
        response = client.get(
            "/api/knowledge/retrieval-catalog/datasets",
            params={"agent_name": "researcher", "page": 1, "page_size": 20},
        )

    assert response.status_code == 200
    assert response.json() == {
        "items": [
            {"id": "dataset-1", "name": "Policies", "selectable": True},
            {"id": "dataset-2", "name": "Empty", "selectable": False},
        ],
        "page": 1,
        "page_size": 20,
        "total": 2,
    }
    assert [call.kwargs["dataset_id"] for call in ragflow.list_datasets.await_args_list] == ["dataset-1", "dataset-2"]


def test_retrieval_catalog_accepts_main_assistant(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _enable_scope_catalog(monkeypatch)
    ragflow = SimpleNamespace(
        list_datasets=AsyncMock(
            return_value=[
                {
                    "id": "dataset-1",
                    "name": "Policies",
                    "embedding_model": "embed-a",
                    "chunk_count": 3,
                }
            ]
        )
    )
    config = _config(scope_selection_enabled=True, datasets=["dataset-1"])

    with TestClient(_app(monkeypatch, ragflow, config=config)) as client:
        response = client.get(
            "/api/knowledge/retrieval-catalog/datasets",
            params={"agent_name": "lead_agent"},
        )

    assert response.status_code == 200
    assert response.json()["items"] == [{"id": "dataset-1", "name": "Policies", "selectable": True}]


def test_retrieval_catalog_documents_reject_outside_allowlist_without_provider_call(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _enable_scope_catalog(monkeypatch)
    ragflow = SimpleNamespace(
        list_datasets=AsyncMock(),
        list_documents=AsyncMock(),
    )
    config = _config(scope_selection_enabled=True, datasets=["dataset-1"])

    with TestClient(_app(monkeypatch, ragflow, config=config)) as client:
        response = client.get(
            "/api/knowledge/retrieval-catalog/datasets/dataset-2/documents",
            params={"agent_name": "researcher"},
        )

    assert response.status_code == 404
    ragflow.list_datasets.assert_not_awaited()
    ragflow.list_documents.assert_not_awaited()


def test_retrieval_catalog_documents_marks_only_searchable_files_selectable(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _enable_scope_catalog(monkeypatch)
    ragflow = SimpleNamespace(
        list_datasets=AsyncMock(
            return_value=[
                {
                    "id": "dataset-1",
                    "name": "Policies",
                    "embedding_model": "embed-a",
                    "chunk_count": 3,
                }
            ]
        ),
        list_documents=AsyncMock(
            return_value={
                "code": 0,
                "data": {
                    "total": 3,
                    "docs": [
                        {
                            "id": "doc-1",
                            "name": "Ready.pdf",
                            "run": "DONE",
                            "chunk_count": 2,
                        },
                        {
                            "id": "doc-2",
                            "name": "Parsing.pdf",
                            "run": "RUNNING",
                            "chunk_count": 0,
                        },
                        {
                            "id": "doc-3",
                            "name": "Empty.pdf",
                            "run": "DONE",
                            "chunk_count": 0,
                        },
                    ],
                },
            }
        ),
    )
    config = _config(scope_selection_enabled=True, datasets=["dataset-1"])

    with TestClient(_app(monkeypatch, ragflow, config=config)) as client:
        response = client.get(
            "/api/knowledge/retrieval-catalog/datasets/dataset-1/documents",
            params={
                "agent_name": "researcher",
                "search": "ready",
                "page": 2,
                "page_size": 10,
            },
        )

    assert response.status_code == 200
    assert response.json()["items"] == [
        {"id": "doc-1", "name": "Ready.pdf", "selectable": True},
        {"id": "doc-2", "name": "Parsing.pdf", "selectable": False},
        {"id": "doc-3", "name": "Empty.pdf", "selectable": False},
    ]
    ragflow.list_documents.assert_awaited_once_with(
        "dataset-1",
        params=[("page", "2"), ("page_size", "10"), ("keywords", "ready")],
    )


@pytest.mark.parametrize(
    "config",
    [
        _config(enabled=False, scope_selection_enabled=True),
        _config(scope_selection_enabled=False),
        _config(
            scope_selection_enabled=True,
            provider=("deerflow.community.lightrag.tools:knowledge_search_tool"),
        ),
    ],
)
def test_retrieval_catalog_fails_closed_when_capability_is_unavailable(
    monkeypatch: pytest.MonkeyPatch,
    config: SimpleNamespace,
) -> None:
    _enable_scope_catalog(monkeypatch)
    ragflow = SimpleNamespace(list_datasets=AsyncMock())

    with TestClient(_app(monkeypatch, ragflow, config=config)) as client:
        response = client.get(
            "/api/knowledge/retrieval-catalog/datasets",
            params={"agent_name": "researcher"},
        )

    assert response.status_code == 409
    ragflow.list_datasets.assert_not_awaited()


def test_management_routes_are_not_exposed(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    with TestClient(_app(monkeypatch, SimpleNamespace())) as client:
        assert client.get("/api/knowledge/datasets").status_code == 404
        assert (
            client.post(
                "/api/knowledge/datasets",
                json={"name": "Deferred"},
            ).status_code
            == 404
        )
        assert client.get("/api/knowledge/events").status_code == 404