diff --git a/test/testcases/conftest.py b/test/testcases/conftest.py index 27826e125c..69c6a394bb 100644 --- a/test/testcases/conftest.py +++ b/test/testcases/conftest.py @@ -14,6 +14,83 @@ # limitations under the License. # +import importlib +import sys +import types + + +def _make_stub_getattr(module_name): + def __getattr__(attr_name): + message = f"{module_name}.{attr_name} is stubbed in tests" + + class _Stub: + def __init__(self, *_args, **_kwargs): + raise RuntimeError(message) + + def __call__(self, *_args, **_kwargs): + raise RuntimeError(message) + + def __getattr__(self, _name): + raise RuntimeError(message) + + setattr(sys.modules[module_name], attr_name, _Stub) + return _Stub + + return __getattr__ + + +def _install_rag_llm_stubs(): + rag_llm = sys.modules.get("rag.llm") + if rag_llm is not None and getattr(rag_llm, "_rag_llm_stubbed", False): + return + + try: + rag_pkg = importlib.import_module("rag") + except Exception: + rag_pkg = types.ModuleType("rag") + rag_pkg.__path__ = [] + rag_pkg.__package__ = "rag" + rag_pkg.__file__ = __file__ + sys.modules["rag"] = rag_pkg + + llm_pkg = types.ModuleType("rag.llm") + llm_pkg.__path__ = [] + llm_pkg.__package__ = "rag.llm" + llm_pkg.__file__ = __file__ + sys.modules["rag.llm"] = llm_pkg + rag_pkg.llm = llm_pkg + + llm_pkg.__getattr__ = _make_stub_getattr("rag.llm") + + for submodule in ("cv_model", "chat_model"): + full_name = f"rag.llm.{submodule}" + sub_mod = sys.modules.get(full_name) + if sub_mod is None or not isinstance(sub_mod, types.ModuleType): + sub_mod = types.ModuleType(full_name) + sys.modules[full_name] = sub_mod + sub_mod.__package__ = "rag.llm" + sub_mod.__file__ = __file__ + sub_mod.__getattr__ = _make_stub_getattr(full_name) + setattr(llm_pkg, submodule, sub_mod) + + llm_pkg._rag_llm_stubbed = True + + +def _install_scholarly_stub(): + if "scholarly" in sys.modules: + return + stub = types.ModuleType("scholarly") + + def _stub(*_args, **_kwargs): + raise RuntimeError("scholarly is stubbed in tests") + + stub.scholarly = _stub + sys.modules["scholarly"] = stub + + +_install_rag_llm_stubs() +_install_scholarly_stub() + import pytest import requests from configs import EMAIL, HOST_ADDRESS, PASSWORD, VERSION, ZHIPU_AI_API_KEY diff --git a/test/testcases/test_http_api/common.py b/test/testcases/test_http_api/common.py index c1567f5742..4e27f74a1e 100644 --- a/test/testcases/test_http_api/common.py +++ b/test/testcases/test_http_api/common.py @@ -304,6 +304,12 @@ def metadata_summary(auth, dataset_id, params=None): return res.json() +def metadata_batch_update(auth, dataset_id, payload=None): + url = f"{HOST_ADDRESS}{DATASETS_API_URL}/{dataset_id}/metadata/update" + res = requests.post(url=url, headers=HEADERS, auth=auth, json=payload) + return res.json() + + # CHAT COMPLETIONS AND RELATED QUESTIONS def related_questions(auth, payload=None): url = f"{HOST_ADDRESS}/api/{VERSION}/sessions/related_questions" diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py new file mode 100644 index 0000000000..63fef2de89 --- /dev/null +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py @@ -0,0 +1,350 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import sys +from copy import deepcopy +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _DummyKB: + def __init__(self, embd_id="embd@factory", chunk_num=1): + self.embd_id = embd_id + self.chunk_num = chunk_num + + def to_json(self): + return {"id": "kb-1"} + + +class _DummyDialogRecord: + def __init__(self): + self._data = { + "id": "chat-1", + "name": "chat-name", + "prompt_config": { + "system": "Answer with {knowledge}", + "parameters": [{"key": "knowledge", "optional": False}], + "prologue": "hello", + "quote": True, + }, + "llm_setting": {"temperature": 0.1}, + "llm_id": "glm-4", + "similarity_threshold": 0.2, + "vector_similarity_weight": 0.3, + "top_n": 6, + "rerank_id": "", + "top_k": 1024, + "kb_ids": ["kb-1"], + "icon": "icon.png", + } + + def to_json(self): + return deepcopy(self._data) + + +def _run(coro): + return asyncio.run(coro) + + +def _load_chat_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + deepdoc_pkg = ModuleType("deepdoc") + deepdoc_parser_pkg = ModuleType("deepdoc.parser") + deepdoc_parser_pkg.__path__ = [] + + class _StubPdfParser: + pass + + class _StubExcelParser: + pass + + class _StubDocxParser: + pass + + deepdoc_parser_pkg.PdfParser = _StubPdfParser + deepdoc_parser_pkg.ExcelParser = _StubExcelParser + deepdoc_parser_pkg.DocxParser = _StubDocxParser + deepdoc_pkg.parser = deepdoc_parser_pkg + monkeypatch.setitem(sys.modules, "deepdoc", deepdoc_pkg) + monkeypatch.setitem(sys.modules, "deepdoc.parser", deepdoc_parser_pkg) + + deepdoc_excel_module = ModuleType("deepdoc.parser.excel_parser") + deepdoc_excel_module.RAGFlowExcelParser = _StubExcelParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.excel_parser", deepdoc_excel_module) + + deepdoc_parser_utils = ModuleType("deepdoc.parser.utils") + deepdoc_parser_utils.get_text = lambda *_args, **_kwargs: "" + monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) + monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + + module_name = "test_chat_sdk_routes_unit_module" + module_path = repo_root / "api" / "apps" / "sdk" / "chat.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _set_request_json(monkeypatch, module, payload): + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(deepcopy(payload))) + + +@pytest.mark.p2 +def test_create_internal_failure_paths(monkeypatch): + module = _load_chat_module(monkeypatch) + + _set_request_json(monkeypatch, module, {"name": "chat-a", "dataset_ids": ["kb-1", "kb-2"]}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: [SimpleNamespace(id="kb")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [_DummyKB(chunk_num=1)]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [_DummyKB(embd_id="embd-a@x"), _DummyKB(embd_id="embd-b@y")]) + monkeypatch.setattr(module.TenantLLMService, "split_model_name_and_factory", lambda model: (model.split("@")[0], "factory")) + res = _run(module.create.__wrapped__("tenant-1")) + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR + assert "different embedding models" in res["message"] + + _set_request_json(monkeypatch, module, {"name": "chat-a", "dataset_ids": []}) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (False, None)) + res = _run(module.create.__wrapped__("tenant-1")) + assert res["message"] == "Tenant not found!" + + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4"))) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module.DialogService, "save", lambda **_kwargs: False) + res = _run(module.create.__wrapped__("tenant-1")) + assert res["message"] == "Fail to new a chat!" + + monkeypatch.setattr(module.DialogService, "save", lambda **_kwargs: True) + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (False, None)) + res = _run(module.create.__wrapped__("tenant-1")) + assert res["message"] == "Fail to new a chat!" + + _set_request_json( + monkeypatch, + module, + {"name": "chat-rerank", "dataset_ids": [], "prompt": {"rerank_model": "unknown-rerank-model"}}, + ) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4"))) + rerank_query_calls = [] + + def _mock_tenant_llm_query(**kwargs): + rerank_query_calls.append(kwargs) + return False + + monkeypatch.setattr(module.TenantLLMService, "query", _mock_tenant_llm_query) + res = _run(module.create.__wrapped__("tenant-1")) + assert "`rerank_model` unknown-rerank-model doesn't exist" in res["message"] + assert rerank_query_calls[-1]["model_type"] == "rerank" + assert rerank_query_calls[-1]["llm_name"] == "unknown-rerank-model" + + _set_request_json(monkeypatch, module, {"name": "chat-tenant", "dataset_ids": [], "tenant_id": "tenant-forbidden"}) + res = _run(module.create.__wrapped__("tenant-1")) + assert res["message"] == "`tenant_id` must not be provided." + + +@pytest.mark.p2 +def test_update_internal_failure_paths(monkeypatch): + module = _load_chat_module(monkeypatch) + + _set_request_json(monkeypatch, module, {"name": "anything"}) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["message"] == "You do not own the chat" + + _set_request_json(monkeypatch, module, {"name": "chat-name"}) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="chat-1")]) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (False, None)) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["message"] == "Tenant not found!" + + _set_request_json(monkeypatch, module, {"dataset_ids": ["kb-1", "kb-2"]}) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(id="tenant-1"))) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: [SimpleNamespace(id="kb")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [_DummyKB(chunk_num=1)]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [_DummyKB(embd_id="embd-a@x"), _DummyKB(embd_id="embd-b@y")]) + monkeypatch.setattr(module.TenantLLMService, "split_model_name_and_factory", lambda model: (model.split("@")[0], "factory")) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR + assert "different embedding models" in res["message"] + + _set_request_json(monkeypatch, module, {"avatar": "new-avatar"}) + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (True, _DummyDialogRecord())) + monkeypatch.setattr(module.DialogService, "update_by_id", lambda *_args, **_kwargs: False) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["message"] == "Chat not found!" + + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(id="tenant-1"))) + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (True, _DummyDialogRecord())) + monkeypatch.setattr(module.DialogService, "update_by_id", lambda *_args, **_kwargs: True) + monkeypatch.setattr( + module.DialogService, + "query", + lambda **kwargs: ( + [SimpleNamespace(id="chat-1")] + if kwargs.get("id") == "chat-1" + else ([SimpleNamespace(id="dup")] if kwargs.get("name") == "dup-name" else []) + ), + ) + monkeypatch.setattr( + module.TenantLLMService, + "split_model_name_and_factory", + lambda model: (model.split("@")[0], "factory"), + ) + monkeypatch.setattr( + module.TenantLLMService, + "query", + lambda **kwargs: kwargs.get("llm_name") in {"glm-4", "allowed-rerank"}, + ) + + _set_request_json(monkeypatch, module, {"show_quotation": True}) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["code"] == 0 + + _set_request_json(monkeypatch, module, {"dataset_ids": ["kb-no-owner"]}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: []) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert "You don't own the dataset kb-no-owner" in res["message"] + + _set_request_json(monkeypatch, module, {"dataset_ids": ["kb-unparsed"]}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: [SimpleNamespace(id="kb-unparsed")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [_DummyKB(chunk_num=0)]) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert "doesn't own parsed file" in res["message"] + + _set_request_json(monkeypatch, module, {"llm": {"model_name": "unknown-model", "model_type": "unsupported"}}) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert "`model_name` unknown-model doesn't exist" in res["message"] + + _set_request_json( + monkeypatch, + module, + {"prompt": {"prompt": "No placeholder", "variables": [{"key": "knowledge", "optional": False}], "rerank_model": "unknown-rerank"}}, + ) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert "`rerank_model` unknown-rerank doesn't exist" in res["message"] + + _set_request_json( + monkeypatch, + module, + {"prompt": {"prompt": "No placeholder", "variables": [{"key": "knowledge", "optional": False}]}}, + ) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert "Parameter 'knowledge' is not used" in res["message"] + + _set_request_json( + monkeypatch, + module, + {"prompt": {"prompt": "Optional-only prompt", "variables": [{"key": "maybe", "optional": True}]}}, + ) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["code"] == 0 + + _set_request_json(monkeypatch, module, {"name": ""}) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["message"] == "`name` cannot be empty." + + _set_request_json(monkeypatch, module, {"name": "dup-name"}) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["message"] == "Duplicated chat name in updating chat." + + _set_request_json(monkeypatch, module, {"llm": {"model_name": "glm-4", "temperature": 0.9}}) + res = _run(module.update.__wrapped__("tenant-1", "chat-1")) + assert res["code"] == 0 + + +@pytest.mark.p2 +def test_delete_duplicate_no_success_path(monkeypatch): + module = _load_chat_module(monkeypatch) + + _set_request_json(monkeypatch, module, {"ids": ["chat-1", "chat-1"]}) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="chat-1")]) + monkeypatch.setattr(module.DialogService, "update_by_id", lambda *_args, **_kwargs: 0) + res = _run(module.delete_chats.__wrapped__("tenant-1")) + assert res["code"] == module.RetCode.DATA_ERROR + assert "Duplicate assistant ids: chat-1" in res["message"] + + _set_request_json(monkeypatch, module, {"ids": ["missing-chat"]}) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) + res = _run(module.delete_chats.__wrapped__("tenant-1")) + assert res["code"] == module.RetCode.DATA_ERROR + assert "Assistant(missing-chat) not found." in res["message"] + + _set_request_json(monkeypatch, module, {"ids": ["chat-1", "chat-1"]}) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="chat-1")]) + monkeypatch.setattr(module.DialogService, "update_by_id", lambda *_args, **_kwargs: 1) + res = _run(module.delete_chats.__wrapped__("tenant-1")) + assert res["code"] == 0 + assert res["data"]["success_count"] == 1 + + +@pytest.mark.p2 +def test_list_missing_kb_warning_and_desc_false(monkeypatch, caplog): + module = _load_chat_module(monkeypatch) + + monkeypatch.setattr(module, "request", SimpleNamespace(args={"desc": "False"})) + monkeypatch.setattr(module.DialogService, "get_list", lambda *_args, **_kwargs: [ + { + "id": "chat-1", + "name": "chat-name", + "prompt_config": {"system": "Answer with {knowledge}", "parameters": [{"key": "knowledge", "optional": False}], "do_refer": True}, + "similarity_threshold": 0.2, + "vector_similarity_weight": 0.3, + "top_n": 6, + "rerank_id": "", + "llm_setting": {"temperature": 0.1}, + "llm_id": "glm-4", + "kb_ids": ["missing-kb"], + "icon": "icon.png", + } + ]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: []) + + with caplog.at_level("WARNING"): + res = module.list_chat.__wrapped__("tenant-1") + + assert res["code"] == 0 + assert res["data"][0]["datasets"] == [] + assert res["data"][0]["avatar"] == "icon.png" + assert "does not exist" in caplog.text diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_create_chat_assistant.py b/test/testcases/test_http_api/test_chat_assistant_management/test_create_chat_assistant.py index 7a588722c9..47a5de905e 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_create_chat_assistant.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_create_chat_assistant.py @@ -231,6 +231,22 @@ class TestChatAssistantCreate: else: assert res["message"] == expected_message + @pytest.mark.p2 + def test_create_additional_guards_p2(self, HttpApiAuth): + tenant_payload = {"name": "guard-tenant-id", "dataset_ids": [], "tenant_id": "tenant-should-not-pass"} + res = create_chat_assistant(HttpApiAuth, tenant_payload) + assert res["code"] == 102 + assert res["message"] == "`tenant_id` must not be provided." + + rerank_payload = { + "name": "guard-rerank-id", + "dataset_ids": [], + "prompt": {"rerank_model": "unknown-rerank-model"}, + } + res = create_chat_assistant(HttpApiAuth, rerank_payload) + assert res["code"] == 102 + assert "`rerank_model` unknown-rerank-model doesn't exist" in res["message"] + class TestChatAssistantCreate2: @pytest.mark.p2 diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_delete_chat_assistants.py b/test/testcases/test_http_api/test_chat_assistant_management/test_delete_chat_assistants.py index 2a2fdc9a6a..670ab04d34 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_delete_chat_assistants.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_delete_chat_assistants.py @@ -125,3 +125,20 @@ class TestChatAssistantsDelete: res = list_chat_assistants(HttpApiAuth) assert len(res["data"]) == 0 + + @pytest.mark.p2 + def test_delete_all_errors_no_success_p2(self, HttpApiAuth, add_chat_assistants_func): + delete_payload = {"ids": ["missing-1", "missing-2"]} + res = delete_chat_assistants(HttpApiAuth, delete_payload) + assert res["code"] == 102 + assert "Assistant(missing-1) not found." in res["message"] + assert "Assistant(missing-2) not found." in res["message"] + + @pytest.mark.p2 + def test_delete_duplicate_partial_success_p2(self, HttpApiAuth, add_chat_assistants_func): + _, _, chat_assistant_ids = add_chat_assistants_func + payload = {"ids": [chat_assistant_ids[0], chat_assistant_ids[0]]} + res = delete_chat_assistants(HttpApiAuth, payload) + assert res["code"] == 0 + assert res["data"]["success_count"] == 1 + assert "Duplicate assistant ids" in res["data"]["errors"][0] diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_list_chat_assistants.py b/test/testcases/test_http_api/test_chat_assistant_management/test_list_chat_assistants.py index 20bce689ee..d9a4697a45 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_list_chat_assistants.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_list_chat_assistants.py @@ -312,3 +312,9 @@ class TestChatAssistantsList: res = list_chat_assistants(HttpApiAuth) assert res["code"] == 0 assert len(res["data"]) == 5 + + @pytest.mark.p2 + def test_desc_false_parse_branch_p2(self, HttpApiAuth): + res = list_chat_assistants(HttpApiAuth, params={"desc": "False", "orderby": "create_time"}) + assert res["code"] == 0 + assert is_sorted(res["data"], "create_time", False) diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_update_chat_assistant.py b/test/testcases/test_http_api/test_chat_assistant_management/test_update_chat_assistant.py index d576821c1a..dbd5d01612 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_update_chat_assistant.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_update_chat_assistant.py @@ -14,7 +14,7 @@ # limitations under the License. # import pytest -from common import list_chat_assistants, update_chat_assistant +from common import create_chat_assistant, list_chat_assistants, update_chat_assistant from configs import CHAT_ASSISTANT_NAME_LIMIT, INVALID_API_TOKEN from libs.auth import RAGFlowHttpApiAuth from utils import encode_avatar @@ -229,3 +229,64 @@ class TestChatAssistantUpdate: ) else: assert expected_message in res["message"] + + @pytest.mark.p2 + def test_update_mapping_and_validation_branches_p2(self, HttpApiAuth, add_chat_assistants_func, chat_assistant_llm_model_type): + dataset_id, _, chat_assistant_ids = add_chat_assistants_func + chat_id = chat_assistant_ids[0] + + res = update_chat_assistant(HttpApiAuth, "invalid-chat-id", {"name": "anything"}) + assert res["code"] == 102 + assert res["message"] == "You do not own the chat" + + res = update_chat_assistant(HttpApiAuth, chat_id, {"show_quotation": False, "dataset_ids": [dataset_id]}) + assert res["code"] == 0 + + res = update_chat_assistant( + HttpApiAuth, + chat_id, + {"llm": {"model_name": "unknown-llm-model", "model_type": chat_assistant_llm_model_type}}, + ) + assert res["code"] == 102 + assert "`model_name` unknown-llm-model doesn't exist" in res["message"] + + res = update_chat_assistant( + HttpApiAuth, + chat_id, + {"prompt": {"rerank_model": "unknown-rerank-model"}}, + ) + assert res["code"] == 102 + assert "`rerank_model` unknown-rerank-model doesn't exist" in res["message"] + + res = update_chat_assistant(HttpApiAuth, chat_id, {"name": ""}) + assert res["code"] == 102 + assert res["message"] == "`name` cannot be empty." + + res = update_chat_assistant(HttpApiAuth, chat_id, {"name": "test_chat_assistant_1"}) + assert res["code"] == 102 + assert res["message"] == "Duplicated chat name in updating chat." + + res = update_chat_assistant( + HttpApiAuth, + chat_id, + {"prompt": {"prompt": "No required placeholder", "variables": [{"key": "knowledge", "optional": False}]}}, + ) + assert res["code"] == 102 + assert "Parameter 'knowledge' is not used" in res["message"] + + res = update_chat_assistant(HttpApiAuth, chat_id, {"avatar": "raw-avatar-value"}) + assert res["code"] == 0 + listed = list_chat_assistants(HttpApiAuth, {"id": chat_id}) + assert listed["code"] == 0 + assert listed["data"][0]["avatar"] == "raw-avatar-value" + + @pytest.mark.p2 + def test_update_unparsed_dataset_guard_p2(self, HttpApiAuth, add_dataset_func, clear_chat_assistants): + dataset_id = add_dataset_func + create_res = create_chat_assistant(HttpApiAuth, {"name": "update-unparsed-target", "dataset_ids": []}) + assert create_res["code"] == 0 + + chat_id = create_res["data"]["id"] + res = update_chat_assistant(HttpApiAuth, chat_id, {"dataset_ids": [dataset_id]}) + assert res["code"] == 102 + assert "doesn't own parsed file" in res["message"] diff --git a/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py b/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py new file mode 100644 index 0000000000..38ec8d9fa0 --- /dev/null +++ b/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py @@ -0,0 +1,219 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import inspect +import sys +from copy import deepcopy +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _DummyKB: + def __init__(self, tenant_id="tenant-1", embd_id="embd-1"): + self.tenant_id = tenant_id + self.embd_id = embd_id + + +class _DummyRetriever: + async def retrieval(self, *_args, **_kwargs): + return { + "chunks": [ + {"doc_id": "doc-1", "content_with_weight": "chunk-content", "similarity": 0.8, "docnm_kwd": "doc-title", "vector": [0.1]} + ] + } + + def retrieval_by_children(self, chunks, _tenant_ids): + return chunks + + +def _run(coro): + return asyncio.run(coro) + + +def _load_dify_retrieval_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + deepdoc_pkg = ModuleType("deepdoc") + deepdoc_parser_pkg = ModuleType("deepdoc.parser") + deepdoc_parser_pkg.__path__ = [] + + class _StubPdfParser: + pass + + class _StubExcelParser: + pass + + class _StubDocxParser: + pass + + deepdoc_parser_pkg.PdfParser = _StubPdfParser + deepdoc_parser_pkg.ExcelParser = _StubExcelParser + deepdoc_parser_pkg.DocxParser = _StubDocxParser + deepdoc_pkg.parser = deepdoc_parser_pkg + monkeypatch.setitem(sys.modules, "deepdoc", deepdoc_pkg) + monkeypatch.setitem(sys.modules, "deepdoc.parser", deepdoc_parser_pkg) + + deepdoc_excel_module = ModuleType("deepdoc.parser.excel_parser") + deepdoc_excel_module.RAGFlowExcelParser = _StubExcelParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.excel_parser", deepdoc_excel_module) + + deepdoc_parser_utils = ModuleType("deepdoc.parser.utils") + deepdoc_parser_utils.get_text = lambda *_args, **_kwargs: "" + monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) + monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + + module_name = "test_dify_retrieval_routes_unit_module" + module_path = repo_root / "api" / "apps" / "sdk" / "dify_retrieval.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _set_request_json(monkeypatch, module, payload): + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(deepcopy(payload))) + + +@pytest.mark.p2 +def test_retrieval_success_with_metadata_and_kg(monkeypatch): + module = _load_dify_retrieval_module(monkeypatch) + _set_request_json( + monkeypatch, + module, + { + "knowledge_id": "kb-1", + "query": "hello", + "use_kg": True, + "retrieval_setting": {"score_threshold": 0.1, "top_k": 3}, + "metadata_condition": {"conditions": [{"name": "author", "comparison_operator": "is", "value": "alice"}], "logic": "and"}, + }, + ) + + monkeypatch.setattr(module, "jsonify", lambda payload: payload) + monkeypatch.setattr(module.DocMetadataService, "get_meta_by_kbs", lambda _kb_ids: [{"doc_id": "doc-1"}]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _DummyKB())) + monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: object()) + monkeypatch.setattr(module, "convert_conditions", lambda cond: cond.get("conditions", [])) + monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: []) + + retriever = _DummyRetriever() + monkeypatch.setattr(module.settings, "retriever", retriever) + + class _DummyKgRetriever: + async def retrieval(self, *_args, **_kwargs): + return { + "doc_id": "doc-2", + "content_with_weight": "kg-content", + "similarity": 0.9, + "docnm_kwd": "kg-title", + } + + monkeypatch.setattr(module.settings, "kg_retriever", _DummyKgRetriever()) + monkeypatch.setattr( + module.DocumentService, + "get_by_id", + lambda doc_id: (True, SimpleNamespace(meta_fields={"origin": f"meta-{doc_id}"})), + ) + monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: []) + + res = _run(inspect.unwrap(module.retrieval)("tenant-1")) + assert "records" in res, res + assert len(res["records"]) == 2, res + top = res["records"][0] + assert top["title"] == "kg-title", res + assert top["metadata"]["doc_id"] == "doc-2", res + assert "score" in top, res + + +@pytest.mark.p2 +def test_retrieval_kb_not_found(monkeypatch): + module = _load_dify_retrieval_module(monkeypatch) + _set_request_json(monkeypatch, module, {"knowledge_id": "kb-missing", "query": "hello"}) + monkeypatch.setattr(module.DocMetadataService, "get_meta_by_kbs", lambda _kb_ids: []) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + + res = _run(inspect.unwrap(module.retrieval)("tenant-1")) + assert res["code"] == module.RetCode.NOT_FOUND, res + assert "Knowledgebase not found" in res["message"], res + + +@pytest.mark.p2 +def test_retrieval_not_found_exception_mapping(monkeypatch): + module = _load_dify_retrieval_module(monkeypatch) + _set_request_json(monkeypatch, module, {"knowledge_id": "kb-1", "query": "hello"}) + monkeypatch.setattr(module.DocMetadataService, "get_meta_by_kbs", lambda _kb_ids: []) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _DummyKB())) + monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: object()) + monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: []) + + class _BrokenRetriever: + async def retrieval(self, *_args, **_kwargs): + raise RuntimeError("chunk_not_found_error") + + monkeypatch.setattr(module.settings, "retriever", _BrokenRetriever()) + + res = _run(inspect.unwrap(module.retrieval)("tenant-1")) + assert res["code"] == module.RetCode.NOT_FOUND, res + assert "No chunk found" in res["message"], res + + +@pytest.mark.p2 +def test_retrieval_generic_exception_mapping(monkeypatch): + module = _load_dify_retrieval_module(monkeypatch) + _set_request_json(monkeypatch, module, {"knowledge_id": "kb-1", "query": "hello"}) + monkeypatch.setattr(module.DocMetadataService, "get_meta_by_kbs", lambda _kb_ids: []) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _DummyKB())) + monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: object()) + monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: []) + + class _BrokenRetriever: + async def retrieval(self, *_args, **_kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(module.settings, "retriever", _BrokenRetriever()) + + res = _run(inspect.unwrap(module.retrieval)("tenant-1")) + assert res["code"] == module.RetCode.SERVER_ERROR, res + assert "boom" in res["message"], res diff --git a/test/testcases/test_http_api/test_file_app/test_file_routes.py b/test/testcases/test_http_api/test_file_app/test_file_routes.py new file mode 100644 index 0000000000..7573142100 --- /dev/null +++ b/test/testcases/test_http_api/test_file_app/test_file_routes.py @@ -0,0 +1,1055 @@ +# +# Copyright 2025 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import asyncio +import functools +import importlib.util +import sys +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + return decorator + + +def _load_files_app(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + api_pkg = ModuleType("api") + api_pkg.__path__ = [str(repo_root / "api")] + monkeypatch.setitem(sys.modules, "api", api_pkg) + + apps_pkg = ModuleType("api.apps") + apps_pkg.__path__ = [str(repo_root / "api" / "apps")] + monkeypatch.setitem(sys.modules, "api.apps", apps_pkg) + api_pkg.apps = apps_pkg + + sdk_pkg = ModuleType("api.apps.sdk") + sdk_pkg.__path__ = [str(repo_root / "api" / "apps" / "sdk")] + monkeypatch.setitem(sys.modules, "api.apps.sdk", sdk_pkg) + apps_pkg.sdk = sdk_pkg + + db_pkg = ModuleType("api.db") + db_pkg.__path__ = [] + + class _FileType(Enum): + FOLDER = "folder" + VIRTUAL = "virtual" + DOC = "doc" + VISUAL = "visual" + + db_pkg.FileType = _FileType + monkeypatch.setitem(sys.modules, "api.db", db_pkg) + api_pkg.db = db_pkg + + services_pkg = ModuleType("api.db.services") + services_pkg.__path__ = [] + services_pkg.duplicate_name = lambda _query, **kwargs: kwargs.get("name", "") + monkeypatch.setitem(sys.modules, "api.db.services", services_pkg) + + document_service_mod = ModuleType("api.db.services.document_service") + + class _StubDocumentService: + @staticmethod + def get_by_id(_doc_id): + return True, SimpleNamespace(id=_doc_id) + + @staticmethod + def get_tenant_id(_doc_id): + return "tenant1" + + @staticmethod + def remove_document(*_args, **_kwargs): + return True + + @staticmethod + def update_by_id(*_args, **_kwargs): + return True + + @staticmethod + def insert(_doc): + return SimpleNamespace(id="doc1") + + document_service_mod.DocumentService = _StubDocumentService + monkeypatch.setitem(sys.modules, "api.db.services.document_service", document_service_mod) + services_pkg.document_service = document_service_mod + + file2document_service_mod = ModuleType("api.db.services.file2document_service") + + class _StubFile2DocumentService: + @staticmethod + def get_by_file_id(_file_id): + return [] + + @staticmethod + def delete_by_file_id(*_args, **_kwargs): + return None + + @staticmethod + def get_storage_address(**_kwargs): + return "bucket", "location" + + @staticmethod + def insert(_data): + return SimpleNamespace(to_json=lambda: {}) + + file2document_service_mod.File2DocumentService = _StubFile2DocumentService + monkeypatch.setitem(sys.modules, "api.db.services.file2document_service", file2document_service_mod) + services_pkg.file2document_service = file2document_service_mod + + knowledgebase_service_mod = ModuleType("api.db.services.knowledgebase_service") + + class _StubKnowledgebaseService: + @staticmethod + def get_by_id(_kb_id): + return False, None + + knowledgebase_service_mod.KnowledgebaseService = _StubKnowledgebaseService + monkeypatch.setitem(sys.modules, "api.db.services.knowledgebase_service", knowledgebase_service_mod) + services_pkg.knowledgebase_service = knowledgebase_service_mod + + file_service_mod = ModuleType("api.db.services.file_service") + + class _StubFileService: + @staticmethod + def get_root_folder(_tenant_id): + return {"id": "root"} + + @staticmethod + def get_by_id(_file_id): + return True, SimpleNamespace(id=_file_id, parent_id="root", location="file", tenant_id="tenant1") + + @staticmethod + def get_id_list_by_id(_pf_id, _file_obj_names, _idx, ids): + return ids + + @staticmethod + def create_folder(_file, parent_id, _file_obj_names, _len_id_list): + return SimpleNamespace(id=parent_id) + + @staticmethod + def query(**_kwargs): + return [] + + @staticmethod + def insert(data): + return SimpleNamespace(to_json=lambda: data) + + @staticmethod + def is_parent_folder_exist(_pf_id): + return True + + @staticmethod + def get_by_pf_id(*_args, **_kwargs): + return [], 0 + + @staticmethod + def get_parent_folder(_file_id): + return SimpleNamespace(to_json=lambda: {"id": "root"}) + + @staticmethod + def get_all_parent_folders(_file_id): + return [] + + @staticmethod + def get_all_innermost_file_ids(_file_id, _acc): + return [] + + @staticmethod + def delete_folder_by_pf_id(*_args, **_kwargs): + return None + + @staticmethod + def delete(_file): + return True + + @staticmethod + def update_by_id(*_args, **_kwargs): + return True + + @staticmethod + def get_by_ids(_file_ids): + return [] + + @staticmethod + def move_file(*_args, **_kwargs): + return None + + @staticmethod + def init_knowledgebase_docs(*_args, **_kwargs): + return None + + @staticmethod + def get_parser(_file_type, _file_name, parser_id): + return parser_id + + file_service_mod.FileService = _StubFileService + monkeypatch.setitem(sys.modules, "api.db.services.file_service", file_service_mod) + services_pkg.file_service = file_service_mod + + api_utils_mod = ModuleType("api.utils.api_utils") + + def get_json_result(data=None, message="", code=0): + return {"code": code, "data": data, "message": message} + + async def get_request_json(): + return {} + + def server_error_response(err): + return {"code": 100, "data": None, "message": str(err)} + + def token_required(func): + @functools.wraps(func) + async def _wrapper(*args, **kwargs): + return await func(*args, **kwargs) + + return _wrapper + + api_utils_mod.get_json_result = get_json_result + api_utils_mod.get_request_json = get_request_json + api_utils_mod.server_error_response = server_error_response + api_utils_mod.token_required = token_required + monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod) + + file_utils_mod = ModuleType("api.utils.file_utils") + file_utils_mod.filename_type = lambda _filename: _FileType.DOC.value + monkeypatch.setitem(sys.modules, "api.utils.file_utils", file_utils_mod) + + web_utils_mod = ModuleType("api.utils.web_utils") + web_utils_mod.CONTENT_TYPE_MAP = {"txt": "text/plain", "json": "application/json"} + web_utils_mod.apply_safe_file_response_headers = lambda response, *_args, **_kwargs: response + monkeypatch.setitem(sys.modules, "api.utils.web_utils", web_utils_mod) + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + common_pkg.settings = SimpleNamespace( + STORAGE_IMPL=SimpleNamespace( + obj_exist=lambda *_args, **_kwargs: False, + put=lambda *_args, **_kwargs: None, + get=lambda *_args, **_kwargs: b"", + rm=lambda *_args, **_kwargs: None, + ) + ) + monkeypatch.setitem(sys.modules, "common", common_pkg) + + misc_utils_mod = ModuleType("common.misc_utils") + misc_utils_mod.get_uuid = lambda: "uuid" + + async def thread_pool_exec(func, *args, **kwargs): + return func(*args, **kwargs) + + misc_utils_mod.thread_pool_exec = thread_pool_exec + monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils_mod) + + constants_mod = ModuleType("common.constants") + + class _RetCode: + SUCCESS = 0 + BAD_REQUEST = 400 + NOT_FOUND = 404 + CONFLICT = 409 + SERVER_ERROR = 500 + + constants_mod.RetCode = _RetCode + monkeypatch.setitem(sys.modules, "common.constants", constants_mod) + + module_path = repo_root / "api" / "apps" / "sdk" / "files.py" + spec = importlib.util.spec_from_file_location("api.apps.sdk.files", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, "api.apps.sdk.files", module) + spec.loader.exec_module(module) + return module + + +def _run(coro): + return asyncio.run(coro) + + +class _DummyFile: + def __init__(self, file_id, file_type, name="doc.txt", tenant_id="tenant1", parent_id="parent1", location=None): + self.id = file_id + self.type = file_type + self.name = name + self.location = location or name + self.size = 1 + self.tenant_id = tenant_id + self.parent_id = parent_id + + def to_json(self): + return {"id": self.id, "name": self.name, "type": self.type} + + +class _FalsyFile(_DummyFile): + def __bool__(self): + return False + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _Args(dict): + def get(self, key, default=None, type=None): + value = super().get(key, default) + if value is None or type is None: + return value + try: + return type(value) + except (TypeError, ValueError): + return default + + +class _DummyRequest: + def __init__(self, *, args=None, form=None, files=None): + self.args = _Args(args or {}) + self.form = _AwaitableValue(form or {}) + self.files = _AwaitableValue(files if files is not None else _DummyFiles()) + + +class _DummyUploadFile: + def __init__(self, filename, blob=b"file-bytes"): + self.filename = filename + self._blob = blob + + def read(self): + return self._blob + + +class _DummyFiles(dict): + def __init__(self, file_objs=None): + super().__init__() + self._file_objs = file_objs or [] + if file_objs is not None: + self["file"] = self._file_objs + + def getlist(self, key): + if key == "file": + return list(self._file_objs) + return [] + + +class _DummyResponse: + def __init__(self, data): + self.data = data + self.headers = {} + + +@pytest.mark.p2 +class TestFileMoveUnit: + def test_move_success_and_invalid_parent(self, monkeypatch): + module = _load_files_app(monkeypatch) + file_id = "file1" + parent_id = "parent1" + + async def fake_request_json(): + return {"src_file_ids": [file_id], "dest_file_id": parent_id} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_DummyFile(file_id, module.FileType.DOC.value)]) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _pid: (True, _DummyFile(parent_id, module.FileType.FOLDER.value))) + monkeypatch.setattr(module.FileService, "move_file", lambda *_args, **_kwargs: None) + + res = _run(module.move.__wrapped__("tenant1")) + assert res["code"] == 0 + assert res["data"] is True + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _pid: (False, None)) + res = _run(module.move.__wrapped__("tenant1")) + assert res["code"] == 404 + assert res["message"] == "Parent Folder not found!" + + def test_move_missing_payload(self, monkeypatch): + module = _load_files_app(monkeypatch) + + async def fake_request_json(): + return {} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.move.__wrapped__("tenant1")) + assert res["code"] == 100 + + def test_move_missing_source_branch(self, monkeypatch): + module = _load_files_app(monkeypatch) + + async def fake_request_json(): + return {"src_file_ids": ["file1"], "dest_file_id": "parent1"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_FalsyFile("file1", module.FileType.DOC.value)]) + res = _run(module.move.__wrapped__("tenant1")) + assert res["code"] == 404 + assert res["message"] == "File or Folder not found!" + + +@pytest.mark.p2 +class TestFileConvertUnit: + def test_convert_success_and_delete(self, monkeypatch): + module = _load_files_app(monkeypatch) + file_id = "file1" + kb_id = "kb1" + + async def fake_request_json(): + return {"kb_ids": [kb_id], "file_ids": [file_id]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_DummyFile(file_id, module.FileType.DOC.value)]) + + class _Inform: + document_id = "doc1" + + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda _id: [_Inform()]) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyFile("doc1", module.FileType.DOC.value))) + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant1") + monkeypatch.setattr(module.DocumentService, "remove_document", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.File2DocumentService, "delete_by_file_id", lambda *_args, **_kwargs: None) + + class _Kb: + id = kb_id + parser_id = "parser" + parser_config = {} + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _Kb())) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _id: (True, _DummyFile(file_id, module.FileType.DOC.value))) + + class _Doc: + def __init__(self, doc_id): + self.id = doc_id + + monkeypatch.setattr(module.DocumentService, "insert", lambda _doc: _Doc("newdoc")) + + class _File2Doc: + def to_json(self): + return {"file_id": file_id, "document_id": "newdoc"} + + monkeypatch.setattr(module.File2DocumentService, "insert", lambda _data: _File2Doc()) + + res = _run(module.convert.__wrapped__("tenant1")) + assert res["code"] == 0 + assert res["data"] + + def test_convert_folder(self, monkeypatch): + module = _load_files_app(monkeypatch) + kb_id = "kb1" + + async def fake_request_json(): + return {"kb_ids": [kb_id], "file_ids": ["folder1"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_DummyFile("folder1", module.FileType.FOLDER.value, name="folder")]) + monkeypatch.setattr(module.FileService, "get_all_innermost_file_ids", lambda *_args, **_kwargs: ["inner1"]) + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda _id: []) + monkeypatch.setattr(module.File2DocumentService, "delete_by_file_id", lambda *_args, **_kwargs: None) + + class _Kb: + id = kb_id + parser_id = "parser" + parser_config = {} + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _Kb())) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _id: (True, _DummyFile("inner1", module.FileType.DOC.value))) + monkeypatch.setattr(module.DocumentService, "insert", lambda _doc: _DummyFile("doc1", module.FileType.DOC.value)) + monkeypatch.setattr(module.File2DocumentService, "insert", lambda _data: SimpleNamespace(to_json=lambda: {"file_id": "inner1"})) + + res = _run(module.convert.__wrapped__("tenant1")) + assert res["code"] == 0 + assert res["data"] + + def test_convert_invalid_file_id(self, monkeypatch): + module = _load_files_app(monkeypatch) + + async def fake_request_json(): + return {"kb_ids": ["kb1"], "file_ids": ["missing"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_FalsyFile("missing", module.FileType.DOC.value)]) + res = _run(module.convert.__wrapped__("tenant1")) + assert res["code"] == 404 + assert res["message"] == "File not found!" + + def test_convert_invalid_kb_id(self, monkeypatch): + module = _load_files_app(monkeypatch) + + async def fake_request_json(): + return {"kb_ids": ["missing"], "file_ids": ["file1"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_DummyFile("file1", module.FileType.DOC.value)]) + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda _id: []) + monkeypatch.setattr(module.File2DocumentService, "delete_by_file_id", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + res = _run(module.convert.__wrapped__("tenant1")) + assert res["code"] == 404 + assert res["message"] == "Can't find this dataset!" + + def test_convert_file_missing_second_lookup(self, monkeypatch): + module = _load_files_app(monkeypatch) + + async def fake_request_json(): + return {"kb_ids": ["kb1"], "file_ids": ["file1"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_DummyFile("file1", module.FileType.DOC.value)]) + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda _id: []) + monkeypatch.setattr(module.File2DocumentService, "delete_by_file_id", lambda *_args, **_kwargs: None) + + class _Kb: + id = "kb1" + parser_id = "parser" + parser_config = {} + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _Kb())) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _id: (False, None)) + res = _run(module.convert.__wrapped__("tenant1")) + assert res["code"] == 404 + assert res["message"] == "Can't find this file!" + + def test_convert_missing_payload(self, monkeypatch): + module = _load_files_app(monkeypatch) + + async def fake_request_json(): + return {} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + with pytest.raises(KeyError): + _run(module.convert.__wrapped__("tenant1")) + + +@pytest.mark.p2 +class TestFileRouteBranchUnit: + def test_upload_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _tenant_id: {"id": "root"}) + + # Missing file part. + monkeypatch.setattr(module, "request", _DummyRequest(form={}, files=_DummyFiles())) + res = _run(module.upload.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.BAD_REQUEST + assert res["message"] == "No file part!" + + # Empty filename. + monkeypatch.setattr( + module, + "request", + _DummyRequest(form={"parent_id": "pf1"}, files=_DummyFiles([_DummyUploadFile("")])), + ) + res = _run(module.upload.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.BAD_REQUEST + assert res["message"] == "No selected file!" + + # Parent folder missing. + monkeypatch.setattr( + module, + "request", + _DummyRequest(form={"parent_id": "pf1"}, files=_DummyFiles([_DummyUploadFile("a.txt")])), + ) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (False, None)) + res = _run(module.upload.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Can't find this folder!" + + # Missing folder in branch: file_len != len_id_list. + monkeypatch.setattr( + module, + "request", + _DummyRequest(form={"parent_id": "pf1"}, files=_DummyFiles([_DummyUploadFile("dir/a.txt")])), + ) + monkeypatch.setattr(module.FileService, "get_id_list_by_id", lambda *_args, **_kwargs: ["pf1", "missing-child"]) + + def get_by_id_missing_child(file_id): + if file_id == "missing-child": + return False, None + return True, SimpleNamespace(id="pf1") + + monkeypatch.setattr(module.FileService, "get_by_id", get_by_id_missing_child) + res = _run(module.upload.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Folder not found!" + + # Missing folder in branch: file_len == len_id_list. + monkeypatch.setattr( + module, + "request", + _DummyRequest(form={"parent_id": "pf1"}, files=_DummyFiles([_DummyUploadFile("b.txt")])), + ) + monkeypatch.setattr(module.FileService, "get_id_list_by_id", lambda *_args, **_kwargs: ["pf1", "leaf"]) + pf1_calls = {"count": 0} + + def get_by_id_missing_parent_in_else(file_id): + if file_id == "pf1": + pf1_calls["count"] += 1 + if pf1_calls["count"] == 1: + return True, SimpleNamespace(id="pf1") + return False, None + return True, SimpleNamespace(id=file_id) + + monkeypatch.setattr(module.FileService, "get_by_id", get_by_id_missing_parent_in_else) + res = _run(module.upload.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Folder not found!" + + class _Storage: + def __init__(self): + self.obj_calls = 0 + self.put_calls = [] + + def obj_exist(self, _bucket, _location): + self.obj_calls += 1 + return self.obj_calls == 1 + + def put(self, bucket, location, blob): + self.put_calls.append((bucket, location, blob)) + + storage = _Storage() + monkeypatch.setattr(module.settings, "STORAGE_IMPL", storage) + monkeypatch.setattr( + module, + "request", + _DummyRequest( + form={"parent_id": "pf1"}, + files=_DummyFiles([_DummyUploadFile("dir/a.txt", b"a"), _DummyUploadFile("b.txt", b"b")]), + ), + ) + + def fake_get_by_id(file_id): + if file_id == "mid-id": + return True, SimpleNamespace(id="mid-id") + return True, SimpleNamespace(id="pf1") + + def fake_get_id_list_by_id(_pf_id, file_obj_names, _idx, _ids): + if file_obj_names[-1] == "a.txt": + return ["pf1", "mid-id"] + return ["pf1", "leaf-id"] + + def fake_create_folder(_file, parent_id, _file_obj_names, _len_id_list): + return SimpleNamespace(id=f"{parent_id}-folder") + + monkeypatch.setattr(module.FileService, "get_by_id", fake_get_by_id) + monkeypatch.setattr(module.FileService, "get_id_list_by_id", fake_get_id_list_by_id) + monkeypatch.setattr(module.FileService, "create_folder", fake_create_folder) + monkeypatch.setattr(module, "filename_type", lambda _name: module.FileType.DOC.value) + monkeypatch.setattr(module, "duplicate_name", lambda _query, **kwargs: kwargs["name"]) + monkeypatch.setattr(module, "get_uuid", lambda: "file-id") + monkeypatch.setattr(module.FileService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module.FileService, "insert", lambda data: SimpleNamespace(to_json=lambda: {"id": data["id"]})) + res = _run(module.upload.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SUCCESS + assert len(res["data"]) == 2 + assert storage.put_calls + + # Exception path. + monkeypatch.setattr( + module, + "request", + _DummyRequest(form={"parent_id": "pf1"}, files=_DummyFiles([_DummyUploadFile("boom.txt")])), + ) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (_ for _ in ()).throw(RuntimeError("upload boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.upload.__wrapped__("tenant1")) + assert res["code"] == 500 + assert "upload boom" in res["message"] + + def test_create_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + state = {"req": {"name": "file1"}} + + async def fake_request_json(): + return state["req"] + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _tenant_id: {"id": "root"}) + monkeypatch.setattr(module.FileService, "is_parent_folder_exist", lambda _pf_id: False) + res = _run(module.create.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.BAD_REQUEST + assert "Parent Folder Doesn't Exist!" in res["message"] + + state["req"] = {"name": "dup", "parent_id": "pf1"} + monkeypatch.setattr(module.FileService, "is_parent_folder_exist", lambda _pf_id: True) + monkeypatch.setattr(module.FileService, "query", lambda **_kwargs: [object()]) + res = _run(module.create.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.CONFLICT + assert "Duplicated folder name" in res["message"] + + inserted = {} + + def fake_insert(data): + inserted["payload"] = data + return SimpleNamespace(to_json=lambda: data) + + monkeypatch.setattr(module.FileService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module, "get_uuid", lambda: "uuid-folder") + monkeypatch.setattr(module.FileService, "insert", fake_insert) + + state["req"] = {"name": "folder", "parent_id": "pf1", "type": module.FileType.FOLDER.value} + res = _run(module.create.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SUCCESS + assert inserted["payload"]["type"] == module.FileType.FOLDER.value + + state["req"] = {"name": "virtual", "parent_id": "pf1", "type": "UNKNOWN"} + res = _run(module.create.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SUCCESS + assert inserted["payload"]["type"] == module.FileType.VIRTUAL.value + + monkeypatch.setattr(module.FileService, "is_parent_folder_exist", lambda _pf_id: (_ for _ in ()).throw(RuntimeError("create boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.create.__wrapped__("tenant1")) + assert res["code"] == 500 + assert "create boom" in res["message"] + + def test_list_files_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + calls = {"init": 0} + + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _tenant_id: {"id": "root"}) + monkeypatch.setattr( + module.FileService, + "init_knowledgebase_docs", + lambda _pf_id, _tenant_id: calls.__setitem__("init", calls["init"] + 1), + ) + monkeypatch.setattr(module, "request", _DummyRequest(args={})) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _pf_id: (False, None)) + res = _run(module.list_files.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Folder not found!" + assert calls["init"] == 1 + + monkeypatch.setattr( + module, + "request", + _DummyRequest(args={"parent_id": "p1", "keywords": "k", "page": "2", "page_size": "10", "orderby": "name", "desc": "False"}), + ) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _pf_id: (True, SimpleNamespace(id="p1"))) + monkeypatch.setattr(module.FileService, "get_by_pf_id", lambda *_args, **_kwargs: ([{"id": "f1"}], 1)) + monkeypatch.setattr(module.FileService, "get_parent_folder", lambda _pf_id: None) + res = _run(module.list_files.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "File not found!" + + monkeypatch.setattr(module.FileService, "get_parent_folder", lambda _pf_id: SimpleNamespace(to_json=lambda: {"id": "p0"})) + res = _run(module.list_files.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["total"] == 1 + assert res["data"]["parent_folder"]["id"] == "p0" + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _pf_id: (_ for _ in ()).throw(RuntimeError("list boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.list_files.__wrapped__("tenant1")) + assert res["code"] == 500 + assert "list boom" in res["message"] + + def test_get_root_folder_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _tenant_id: {"id": "root"}) + res = _run(module.get_root_folder.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["root_folder"]["id"] == "root" + + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _tenant_id: (_ for _ in ()).throw(RuntimeError("root boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.get_root_folder.__wrapped__("tenant1")) + assert res["code"] == 500 + assert "root boom" in res["message"] + + def test_get_parent_folder_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + monkeypatch.setattr(module, "request", _DummyRequest(args={"file_id": "missing"})) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (False, None)) + res = _run(module.get_parent_folder.__wrapped__()) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Folder not found!" + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (True, SimpleNamespace(id="f1"))) + monkeypatch.setattr(module.FileService, "get_parent_folder", lambda _file_id: SimpleNamespace(to_json=lambda: {"id": "p1"})) + res = _run(module.get_parent_folder.__wrapped__()) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["parent_folder"]["id"] == "p1" + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (_ for _ in ()).throw(RuntimeError("parent boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.get_parent_folder.__wrapped__()) + assert res["code"] == 500 + assert "parent boom" in res["message"] + + def test_get_all_parent_folders_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + monkeypatch.setattr(module, "request", _DummyRequest(args={"file_id": "missing"})) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (False, None)) + res = _run(module.get_all_parent_folders.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Folder not found!" + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (True, SimpleNamespace(id="f1"))) + monkeypatch.setattr( + module.FileService, + "get_all_parent_folders", + lambda _file_id: [SimpleNamespace(to_json=lambda: {"id": "p1"})], + ) + res = _run(module.get_all_parent_folders.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["parent_folders"] == [{"id": "p1"}] + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (_ for _ in ()).throw(RuntimeError("all parent boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.get_all_parent_folders.__wrapped__("tenant1")) + assert res["code"] == 500 + assert "all parent boom" in res["message"] + + def test_rm_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + req_state = {"file_ids": ["f1"]} + + async def fake_request_json(): + return req_state + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.settings, "STORAGE_IMPL", SimpleNamespace(rm=lambda *_args, **_kwargs: None)) + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (False, None)) + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "File or Folder not found!" + + monkeypatch.setattr( + module.FileService, + "get_by_id", + lambda _file_id: (True, _DummyFile(_file_id, module.FileType.DOC.value, tenant_id=None)), + ) + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Tenant not found!" + + req_state["file_ids"] = ["folder1"] + + def folder_missing_inner(file_id): + if file_id == "folder1": + return True, _DummyFile("folder1", module.FileType.FOLDER.value, parent_id="pf1") + if file_id == "inner1": + return False, None + return False, None + + monkeypatch.setattr(module.FileService, "get_by_id", folder_missing_inner) + monkeypatch.setattr(module.FileService, "get_all_innermost_file_ids", lambda _file_id, _acc: ["inner1"]) + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "File not found!" + + req_state["file_ids"] = ["doc1"] + monkeypatch.setattr( + module.FileService, + "get_by_id", + lambda _file_id: (True, _DummyFile("doc1", module.FileType.DOC.value, parent_id="pf1")), + ) + monkeypatch.setattr(module.FileService, "delete", lambda _file: False) + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SERVER_ERROR + assert "Database error (File removal)!" in res["message"] + + class _Inform: + document_id = "doc1" + + monkeypatch.setattr(module.FileService, "delete", lambda _file: True) + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda _file_id: [_Inform()]) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None)) + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Document not found!" + + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, SimpleNamespace(id=_doc_id))) + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: None) + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Tenant not found!" + + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant1") + monkeypatch.setattr(module.DocumentService, "remove_document", lambda *_args, **_kwargs: False) + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SERVER_ERROR + assert "Database error (Document removal)!" in res["message"] + + req_state["file_ids"] = ["folder-ok"] + deleted = {"folder": 0, "link": 0} + + def folder_success(file_id): + if file_id == "folder-ok": + return True, _DummyFile("folder-ok", module.FileType.FOLDER.value, parent_id="pf1") + if file_id == "inner-ok": + return True, _DummyFile("inner-ok", module.FileType.DOC.value, parent_id="pf1", location="inner.bin") + return False, None + + monkeypatch.setattr(module.FileService, "get_by_id", folder_success) + monkeypatch.setattr(module.FileService, "get_all_innermost_file_ids", lambda _file_id, _acc: ["inner-ok"]) + monkeypatch.setattr( + module.FileService, + "delete_folder_by_pf_id", + lambda _tenant_id, _file_id: deleted.__setitem__("folder", deleted["folder"] + 1), + ) + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda _file_id: []) + monkeypatch.setattr( + module.File2DocumentService, + "delete_by_file_id", + lambda _file_id: deleted.__setitem__("link", deleted["link"] + 1), + ) + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"] is True + assert deleted == {"folder": 1, "link": 1} + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (_ for _ in ()).throw(RuntimeError("rm boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + req_state["file_ids"] = ["boom"] + res = _run(module.rm.__wrapped__("tenant1")) + assert res["code"] == 500 + assert "rm boom" in res["message"] + + def test_rename_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + req_state = {"file_id": "f1", "name": "new.txt"} + + async def fake_request_json(): + return req_state + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (False, None)) + res = _run(module.rename.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "File not found!" + + monkeypatch.setattr( + module.FileService, + "get_by_id", + lambda _file_id: (True, _DummyFile("f1", module.FileType.DOC.value, name="origin.txt")), + ) + req_state["name"] = "new.pdf" + res = _run(module.rename.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.BAD_REQUEST + assert "extension of file can't be changed" in res["message"] + + req_state["name"] = "new.txt" + monkeypatch.setattr(module.FileService, "query", lambda **_kwargs: [SimpleNamespace(name="new.txt")]) + res = _run(module.rename.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.CONFLICT + assert "Duplicated file name in the same folder." in res["message"] + + monkeypatch.setattr(module.FileService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module.FileService, "update_by_id", lambda *_args, **_kwargs: False) + res = _run(module.rename.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SERVER_ERROR + assert "Database error (File rename)!" in res["message"] + + monkeypatch.setattr(module.FileService, "update_by_id", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda _file_id: [SimpleNamespace(document_id="doc1")]) + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: False) + res = _run(module.rename.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SERVER_ERROR + assert "Database error (Document rename)!" in res["message"] + + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda _file_id: []) + res = _run(module.rename.__wrapped__("tenant1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"] is True + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (_ for _ in ()).throw(RuntimeError("rename boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.rename.__wrapped__("tenant1")) + assert res["code"] == 500 + assert "rename boom" in res["message"] + + def test_get_file_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (False, None)) + res = _run(module.get.__wrapped__("tenant1", "missing")) + assert res["code"] == module.RetCode.NOT_FOUND + assert res["message"] == "Document not found!" + + class _Storage: + def __init__(self): + self.calls = 0 + + def get(self, _bucket, _location): + self.calls += 1 + if self.calls == 1: + return None + return b"blob-data" + + storage = _Storage() + monkeypatch.setattr(module.settings, "STORAGE_IMPL", storage) + monkeypatch.setattr( + module.FileService, + "get_by_id", + lambda _file_id: (True, _DummyFile("f1", module.FileType.VISUAL.value, name="image.abc", parent_id="pf1", location="loc1")), + ) + monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("pf2", "loc2")) + + async def fake_make_response(data): + return _DummyResponse(data) + + monkeypatch.setattr(module, "make_response", fake_make_response) + monkeypatch.setattr( + module, + "apply_safe_file_response_headers", + lambda response, content_type, extension: response.headers.update( + {"content_type": content_type, "extension": extension} + ), + ) + res = _run(module.get.__wrapped__("tenant1", "f1")) + assert isinstance(res, _DummyResponse) + assert res.data == b"blob-data" + assert res.headers["extension"] == "abc" + assert res.headers["content_type"] == "image/abc" + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (_ for _ in ()).throw(RuntimeError("get boom"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.get.__wrapped__("tenant1", "f1")) + assert res["code"] == 500 + assert "get boom" in res["message"] + + def test_download_attachment_branch_matrix(self, monkeypatch): + module = _load_files_app(monkeypatch) + monkeypatch.setattr(module, "request", _DummyRequest(args={"ext": "abc"})) + + async def fake_thread_pool_exec(_fn, _tenant_id, _attachment_id): + return b"attachment" + + async def fake_make_response(data): + return _DummyResponse(data) + + monkeypatch.setattr(module, "thread_pool_exec", fake_thread_pool_exec) + monkeypatch.setattr(module, "make_response", fake_make_response) + monkeypatch.setattr( + module, + "apply_safe_file_response_headers", + lambda response, content_type, extension: response.headers.update( + {"content_type": content_type, "extension": extension} + ), + ) + res = _run(module.download_attachment.__wrapped__("tenant1", "att1")) + assert isinstance(res, _DummyResponse) + assert res.data == b"attachment" + assert res.headers["extension"] == "abc" + assert res.headers["content_type"] == "application/abc" diff --git a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py new file mode 100644 index 0000000000..b1fe88f931 --- /dev/null +++ b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py @@ -0,0 +1,992 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import asyncio +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import numpy as np +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _DummyFiles(dict): + def getlist(self, key): + return self.get(key, []) + + +class _DummyArgs(dict): + def getlist(self, key): + v = self.get(key, []) + if v is None: + return [] + if isinstance(v, list): + return v + return [v] + + +class _DummyDoc: + def __init__( + self, + *, + doc_id="doc-1", + kb_id="kb-1", + name="doc.txt", + chunk_num=1, + token_num=2, + progress=0, + process_duration=0, + parser_id="naive", + doc_type=1, + status=True, + run=0, + ): + self.id = doc_id + self.kb_id = kb_id + self.name = name + self.chunk_num = chunk_num + self.token_num = token_num + self.progress = progress + self.process_duration = process_duration + self.parser_id = parser_id + self.type = doc_type + self.status = status + self.run = run + + def to_dict(self): + return { + "id": self.id, + "kb_id": self.kb_id, + "name": self.name, + "chunk_num": self.chunk_num, + "token_num": self.token_num, + "progress": self.progress, + "process_duration": self.process_duration, + "parser_id": self.parser_id, + "run": self.run, + "status": self.status, + } + + +class _ToggleBoolDocList: + def __init__(self, value): + self._calls = 0 + self._value = value + + def __getitem__(self, item): + return self._value + + def __bool__(self): + self._calls += 1 + return self._calls == 1 + + +def _run(coro): + return asyncio.run(coro) + + +def _load_doc_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + deepdoc_pkg = ModuleType("deepdoc") + deepdoc_parser_pkg = ModuleType("deepdoc.parser") + deepdoc_parser_pkg.__path__ = [] + + class _StubPdfParser: + pass + + class _StubExcelParser: + pass + + class _StubDocxParser: + pass + + deepdoc_parser_pkg.PdfParser = _StubPdfParser + deepdoc_parser_pkg.ExcelParser = _StubExcelParser + deepdoc_parser_pkg.DocxParser = _StubDocxParser + deepdoc_pkg.parser = deepdoc_parser_pkg + monkeypatch.setitem(sys.modules, "deepdoc", deepdoc_pkg) + monkeypatch.setitem(sys.modules, "deepdoc.parser", deepdoc_parser_pkg) + + deepdoc_excel_module = ModuleType("deepdoc.parser.excel_parser") + deepdoc_excel_module.RAGFlowExcelParser = _StubExcelParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.excel_parser", deepdoc_excel_module) + deepdoc_parser_utils = ModuleType("deepdoc.parser.utils") + deepdoc_parser_utils.get_text = lambda *_args, **_kwargs: "" + monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) + monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + + module_path = repo_root / "api" / "apps" / "sdk" / "doc.py" + spec = importlib.util.spec_from_file_location("test_doc_sdk_routes_unit", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + spec.loader.exec_module(module) + return module + + +def _patch_send_file(monkeypatch, module): + async def _fake_send_file(file_obj, **kwargs): + return {"file": file_obj, "filename": kwargs.get("attachment_filename")} + + monkeypatch.setattr(module, "send_file", _fake_send_file) + + +def _patch_storage(monkeypatch, module, *, file_stream=b"abc"): + storage = SimpleNamespace(get=lambda *_args, **_kwargs: file_stream, rm=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.settings, "STORAGE_IMPL", storage) + + +def _patch_docstore(monkeypatch, module, **kwargs): + defaults = { + "delete": lambda *_args, **_kwargs: 0, + "update": lambda *_args, **_kwargs: None, + "get": lambda *_args, **_kwargs: {}, + "insert": lambda *_args, **_kwargs: None, + "index_exist": lambda *_args, **_kwargs: False, + } + defaults.update(kwargs) + monkeypatch.setattr(module.settings, "docStoreConn", SimpleNamespace(**defaults)) + + +@pytest.mark.p2 +class TestDocRoutesUnit: + def test_chunk_positions_validation_error(self, monkeypatch): + module = _load_doc_module(monkeypatch) + with pytest.raises(ValueError) as exc_info: + module.Chunk(positions=[[1, 2, 3, 4]]) + assert "length of 5" in str(exc_info.value) + + def test_upload_validation_and_upload_error(self, monkeypatch): + module = _load_doc_module(monkeypatch) + + class _FileObj: + def __init__(self, name): + self.filename = name + + monkeypatch.setattr(module, "request", SimpleNamespace(form=_AwaitableValue({}), files=_AwaitableValue(_DummyFiles({"file": [_FileObj("")]})))) + res = _run(module.upload.__wrapped__("ds-1", "tenant-1")) + assert res["code"] == module.RetCode.ARGUMENT_ERROR + assert res["message"] == "No file selected!" + + long_name = "a" * (module.FILE_NAME_LEN_LIMIT + 1) + monkeypatch.setattr(module, "request", SimpleNamespace(form=_AwaitableValue({}), files=_AwaitableValue(_DummyFiles({"file": [_FileObj(long_name)]})))) + res = _run(module.upload.__wrapped__("ds-1", "tenant-1")) + assert res["code"] == module.RetCode.ARGUMENT_ERROR + assert "bytes or less" in res["message"] + + monkeypatch.setattr(module, "request", SimpleNamespace(form=_AwaitableValue({}), files=_AwaitableValue(_DummyFiles({"file": [_FileObj("ok.txt")]})))) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace())) + monkeypatch.setattr(module.FileService, "upload_document", lambda *_args, **_kwargs: (["upload failed"], [])) + res = _run(module.upload.__wrapped__("ds-1", "tenant-1")) + assert res["code"] == module.RetCode.SERVER_ERROR + assert res["message"] == "upload failed" + + def test_update_doc_guards_and_error_paths(self, monkeypatch): + module = _load_doc_module(monkeypatch) + doc = _DummyDoc() + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "You don't own the dataset." + + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [1]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (False, None)) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "Can't find this dataset!" + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1"))) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "doesn't own the document" in res["message"] + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [doc]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_count": 100})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "chunk_count" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"token_count": 100})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "token_count" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"progress": 100})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "progress" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"meta_fields": []})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "meta_fields must be a dictionary" + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"meta_fields": {"k": "v"}})) + monkeypatch.setattr(module.DocMetadataService, "update_document_metadata", lambda *_args, **_kwargs: False) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "Failed to update metadata" + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"name": "a" * (module.FILE_NAME_LEN_LIMIT + 1)})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["code"] == module.RetCode.ARGUMENT_ERROR + assert "bytes or less" in res["message"] + + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: False) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"name": "new.txt"})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "Document rename" in res["message"] + + def test_update_doc_chunk_method_enabled_and_db_error(self, monkeypatch): + module = _load_doc_module(monkeypatch) + visual_doc = _DummyDoc(parser_id="naive", doc_type=module.FileType.VISUAL) + kb = SimpleNamespace(tenant_id="tenant-1") + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [1]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, kb)) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [visual_doc]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_method": "naive"})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "Not supported yet!" + + doc = _DummyDoc(token_num=2, chunk_num=1, parser_id="naive") + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [doc]) + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: False) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_method": "manual"})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "Document not found!" + + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "update_parser_config", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *_args, **_kwargs: False) + _patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_method": "manual"})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "Document not found!" + + monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *_args, **_kwargs: True) + doc_for_enabled = _DummyDoc(status=False) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [doc_for_enabled]) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _id: (True, doc_for_enabled)) + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: False) + _patch_docstore(monkeypatch, module, update=lambda *_args, **_kwargs: None, delete=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"enabled": True})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "Document update" in res["message"] + + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: True) + _patch_docstore(monkeypatch, module, update=lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("boom")), delete=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["code"] == 500 + assert "boom" in res["message"] + + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _id: (False, None)) + _patch_docstore(monkeypatch, module, update=lambda *_args, **_kwargs: None, delete=lambda *_args, **_kwargs: None) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "Dataset created failed" + + # cover token reset + docStore deletion branch + doc_reset = _DummyDoc(token_num=3, chunk_num=2, parser_id="naive", run=0) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [doc_reset]) + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _id: (True, doc_reset)) + _patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: None, update=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_method": "manual"})) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["code"] == 0 + + def _raise_operational_error(_id): + raise module.OperationalError("db down") + + monkeypatch.setattr(module.DocumentService, "get_by_id", _raise_operational_error) + res = _run(module.update_doc.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "Database operation failed" + + def test_download_and_download_doc_errors(self, monkeypatch): + module = _load_doc_module(monkeypatch) + _patch_send_file(monkeypatch, module) + _patch_storage(monkeypatch, module, file_stream=b"") + res = _run(module.download.__wrapped__("tenant-1", "ds-1", "")) + assert res["message"] == "Specify document_id please." + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: []) + res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "do not own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [1]) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "not own the document" in res["message"] + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()]) + monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) + res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["message"] == "This file is empty." + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"})) + res = _run(module.download_doc("doc-1")) + assert "Authorization is not valid" in res["message"] + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer token"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) + res = _run(module.download_doc("doc-1")) + assert "API key is invalid" in res["message"] + + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace()]) + res = _run(module.download_doc("")) + assert res["message"] == "Specify document_id please." + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + res = _run(module.download_doc("doc-1")) + assert "not own the document" in res["message"] + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()]) + monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) + _patch_storage(monkeypatch, module, file_stream=b"") + res = _run(module.download_doc("doc-1")) + assert res["message"] == "This file is empty." + + _patch_storage(monkeypatch, module, file_stream=b"abc") + res = _run(module.download_doc("doc-1")) + assert res["filename"] == "doc.txt" + + def test_list_docs_metadata_filters(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + monkeypatch.setattr(module, "request", SimpleNamespace(args=_DummyArgs())) + res = module.list_docs.__wrapped__("ds-1", "tenant-1") + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr( + module, + "request", + SimpleNamespace( + args=_DummyArgs( + { + "metadata_condition": "{bad json", + } + ) + ), + ) + res = module.list_docs.__wrapped__("ds-1", "tenant-1") + assert res["message"] == "metadata_condition must be valid JSON." + + monkeypatch.setattr(module, "request", SimpleNamespace(args=_DummyArgs({"metadata_condition": "[1]"}))) + res = module.list_docs.__wrapped__("ds-1", "tenant-1") + assert res["message"] == "metadata_condition must be an object." + + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda _kbs: [{"doc_id": "x"}]) + monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: []) + monkeypatch.setattr(module, "convert_conditions", lambda cond: cond) + monkeypatch.setattr( + module, + "request", + SimpleNamespace(args=_DummyArgs({"metadata_condition": '{"conditions":[{"field":"x","op":"eq","value":"y"}]}'})), + ) + res = module.list_docs.__wrapped__("ds-1", "tenant-1") + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["total"] == 0 + + monkeypatch.setattr( + module.DocumentService, + "get_list", + lambda *_args, **_kwargs: ([{"id": "doc-1", "create_time": 100, "run": "0"}], 1), + ) + monkeypatch.setattr( + module, + "request", + SimpleNamespace( + args=_DummyArgs( + { + "create_time_from": "101", + "create_time_to": "200", + } + ) + ), + ) + res = module.list_docs.__wrapped__("ds-1", "tenant-1") + assert res["code"] == 0 + assert res["data"]["docs"] == [] + + def test_metadata_summary_and_batch_update(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module, "convert_conditions", lambda cond: cond) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"selector": {}})) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.metadata_summary.__wrapped__("ds-1", "tenant-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"doc_ids": ["d1"]})) + monkeypatch.setattr(module.DocMetadataService, "get_metadata_summary", lambda *_args, **_kwargs: {"k": 1}) + res = _run(module.metadata_summary.__wrapped__("ds-1", "tenant-1")) + assert res["code"] == 0 + assert res["data"]["summary"] == {"k": 1} + + monkeypatch.setattr(module.DocMetadataService, "get_metadata_summary", lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("x"))) + monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)}) + res = _run(module.metadata_summary.__wrapped__("ds-1", "tenant-1")) + assert res["code"] == 500 + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"selector": [1]})) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert res["message"] == "selector must be an object." + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"selector": {}, "updates": {"k": "v"}, "deletes": []})) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert res["message"] == "updates and deletes must be lists." + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"selector": {"metadata_condition": [1]}, "updates": [], "deletes": []}), + ) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert res["message"] == "metadata_condition must be an object." + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"selector": {"document_ids": "doc-1"}, "updates": [], "deletes": []}), + ) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert res["message"] == "document_ids must be a list." + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"selector": {}, "updates": [{"key": ""}], "deletes": []}), + ) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert "Each update requires key and value." in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"selector": {}, "updates": [], "deletes": [{"x": "y"}]}), + ) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert "Each delete requires key." in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue( + { + "selector": {"document_ids": ["bad"], "metadata_condition": {"conditions": []}}, + "updates": [{"key": "k", "value": "v"}], + "deletes": [], + } + ), + ) + monkeypatch.setattr(module.KnowledgebaseService, "list_documents_by_ids", lambda _ids: ["doc-1"]) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert "do not belong to dataset" in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue( + { + "selector": {"document_ids": ["doc-1"], "metadata_condition": {"conditions": [{"f": "x"}]}}, + "updates": [{"key": "k", "value": "v"}], + "deletes": [], + } + ), + ) + monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: []) + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda _kbs: []) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert res["code"] == 0 + assert res["data"]["updated"] == 0 + assert res["data"]["matched_docs"] == 0 + + monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: ["doc-1"]) + monkeypatch.setattr(module.DocMetadataService, "batch_update_metadata", lambda *_args, **_kwargs: 1) + res = _run(module.metadata_batch_update.__wrapped__("ds-1", "tenant-1")) + assert res["code"] == 0 + assert res["data"]["updated"] == 1 + assert res["data"]["matched_docs"] == 1 + + def test_delete_branches(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.delete.__wrapped__("tenant-1", "ds-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"ids": ["doc-1"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _tenant: {"id": "pf-1"}) + monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _id: (True, _DummyDoc())) + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _id: None) + res = _run(module.delete.__wrapped__("tenant-1", "ds-1")) + assert res["message"] == "Tenant not found!" + + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _id: "tenant-1") + monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) + monkeypatch.setattr(module.DocumentService, "remove_document", lambda *_args, **_kwargs: False) + res = _run(module.delete.__wrapped__("tenant-1", "ds-1")) + assert "Document removal" in res["message"] + + def _raise_get_by_id(_id): + raise RuntimeError("boom") + + monkeypatch.setattr(module.DocumentService, "get_by_id", _raise_get_by_id) + res = _run(module.delete.__wrapped__("tenant-1", "ds-1")) + assert res["code"] == module.RetCode.SERVER_ERROR + assert "boom" in res["message"] + + monkeypatch.setattr(module, "check_duplicate_ids", lambda _ids, _kind: ([], ["Duplicate document ids: doc-1"])) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _id: (False, None)) + res = _run(module.delete.__wrapped__("tenant-1", "ds-1")) + assert res["code"] == module.RetCode.DATA_ERROR + assert "Duplicate document ids" in res["message"] + + def test_parse_branches(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.parse.__wrapped__("tenant-1", "ds-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + toggle_doc = _ToggleBoolDocList(_DummyDoc(progress=0)) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: toggle_doc) + res = _run(module.parse.__wrapped__("tenant-1", "ds-1")) + assert "don't own the document" in res["message"] + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(progress=0.5)]) + res = _run(module.parse.__wrapped__("tenant-1", "ds-1")) + assert "currently being processed" in res["message"] + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(progress=0)]) + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _id: (True, _DummyDoc())) + monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) + _patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.TaskService, "filter_delete", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "queue_tasks", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, ["Duplicate document ids: doc-1"])) + res = _run(module.parse.__wrapped__("tenant-1", "ds-1")) + assert res["code"] == 0 + assert res["data"]["success_count"] == 1 + assert "Duplicate document ids" in res["data"]["errors"][0] + + monkeypatch.setattr(module, "check_duplicate_ids", lambda _ids, _kind: ([], ["Duplicate document ids: doc-1"])) + res = _run(module.parse.__wrapped__("tenant-1", "ds-1")) + assert res["code"] == module.RetCode.DATA_ERROR + assert "Duplicate document ids" in res["message"] + + def test_stop_parsing_branches(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({})) + res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1")) + assert "`document_ids` is required" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1")) + assert "don't own the document" in res["message"] + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(progress=1)]) + res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1")) + assert "progress at 0 or 1" in res["message"] + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(progress=0.3)]) + monkeypatch.setattr(module, "cancel_all_task_of", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: True) + _patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, ["Duplicate document ids: doc-1"])) + res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1")) + assert res["code"] == 0 + assert res["data"]["success_count"] == 1 + assert "Duplicate document ids" in res["data"]["errors"][0] + + monkeypatch.setattr(module, "check_duplicate_ids", lambda _ids, _kind: ([], ["Duplicate document ids: doc-1"])) + res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1")) + assert res["code"] == module.RetCode.DATA_ERROR + assert "Duplicate document ids" in res["message"] + + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(progress=0.3)]) + res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1")) + assert res["code"] == 0 + + def test_list_chunks_branches(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "don't own the document" in res["message"] + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()]) + monkeypatch.setattr(module, "request", SimpleNamespace(args=_DummyArgs({"id": "chunk-1"}))) + _patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: None) + res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "Chunk not found" in res["message"] + + _patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"id_vec": [1], "content_with_weight_vec": [2]}) + res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "Chunk `chunk-1` not found." in res["message"] + + _patch_docstore( + monkeypatch, + module, + get=lambda *_args, **_kwargs: { + "chunk_id": "chunk-1", + "content_with_weight": "x", + "doc_id": "doc-1", + "docnm_kwd": "doc", + "position_int": [[1, 2, 3, 4, 5]], + }, + ) + res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["code"] == 0 + assert res["data"]["total"] == 1 + assert res["data"]["chunks"][0]["id"] == "chunk-1" + + def test_add_chunk_access_guard(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.add_chunk.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "don't own the dataset" in res["message"] + + def test_rm_chunk_branches(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.rm_chunk.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "get_by_ids", lambda _ids: []) + with pytest.raises(LookupError): + _run(module.rm_chunk.__wrapped__("tenant-1", "ds-1", "doc-1")) + + monkeypatch.setattr(module.DocumentService, "get_by_ids", lambda _ids: [_DummyDoc()]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({})) + _patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: 2) + monkeypatch.setattr(module.DocumentService, "decrement_chunk_num", lambda *_args, **_kwargs: None) + res = _run(module.rm_chunk.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["code"] == 0 + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_ids": ["c1", "c1"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda _ids, _kind: (["c1"], ["Duplicate chunk ids: c1"])) + _patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: 1) + res = _run(module.rm_chunk.__wrapped__("tenant-1", "ds-1", "doc-1")) + assert res["code"] == 0 + assert res["data"]["errors"] == ["Duplicate chunk ids: c1"] + + def test_update_chunk_branches(self, monkeypatch): + module = _load_doc_module(monkeypatch) + _patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: None) + res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1")) + assert "Can't find this chunk" in res["message"] + + _patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"content_with_weight": "q\na"}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1")) + assert "don't own the document" in res["message"] + + doc = _DummyDoc(parser_id="naive") + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [doc]) + monkeypatch.setattr(module.rag_tokenizer, "tokenize", lambda text: text or "") + monkeypatch.setattr(module.rag_tokenizer, "fine_grained_tokenize", lambda text: text or "") + monkeypatch.setattr(module.rag_tokenizer, "is_chinese", lambda _text: False) + monkeypatch.setattr(module.DocumentService, "get_embd_id", lambda _doc_id: "embd") + + class _EmbedModel: + def encode(self, _texts): + return [np.array([0.2, 0.8]), np.array([0.3, 0.7])], 1 + + monkeypatch.setattr(module.TenantLLMService, "model_instance", lambda *_args, **_kwargs: _EmbedModel()) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"positions": "bad"})) + res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1")) + assert "`positions` should be a list" in res["message"] + + _patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"content_with_weight": "x"}, update=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"positions": [[1, 2, 3, 4, 5]]})) + res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1")) + assert res["code"] == 0 + + qa_doc = _DummyDoc(parser_id=module.ParserType.QA) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [qa_doc]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"content": "no-separator"})) + res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1")) + assert "Q&A must be separated" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"content": "Q?\nA!"})) + _patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"content_with_weight": "Q?\nA!"}, update=lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "beAdoc", lambda d, *_args, **_kwargs: d) + res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1")) + assert res["code"] == 0 + + def test_retrieval_validation_matrix(self, monkeypatch): + module = _load_doc_module(monkeypatch) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"dataset_ids": "bad"})) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "`dataset_ids` should be a list" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"dataset_ids": ["ds-1"]})) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "don't own the dataset" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="m1"), SimpleNamespace(embd_id="m2")]) + monkeypatch.setattr(module.TenantLLMService, "split_model_name_and_factory", lambda embd_id: (embd_id, "f")) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "different embedding models" in res["message"] + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="m1", tenant_id="tenant-1")]) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "`question` is required." in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": " "}), + ) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert res["code"] == 0 + assert res["data"]["chunks"] == [] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q", "document_ids": "bad"}), + ) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "`documents` should be a list" in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q", "document_ids": ["not-owned"]}), + ) + monkeypatch.setattr(module.KnowledgebaseService, "list_documents_by_ids", lambda _ids: ["doc-1"]) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "don't own the document" in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q", "metadata_condition": {"logic": "and"}}), + ) + monkeypatch.setattr(module.DocMetadataService, "get_meta_by_kbs", lambda _ids: []) + monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: []) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "code" in res + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q", "highlight": "True"}), + ) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="m1", tenant_id="tenant-1")]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", embd_id="m1"))) + + class _Retriever: + async def retrieval(self, *_args, **_kwargs): + return {"chunks": [], "total": 0} + + def retrieval_by_children(self, chunks, *_args, **_kwargs): + return chunks + + monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: SimpleNamespace()) + monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: {}) + monkeypatch.setattr(module.settings, "retriever", _Retriever()) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert res["code"] == 0 + assert res["data"]["chunks"] == [] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q", "highlight": True}), + ) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert res["code"] == 0 + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q", "highlight": "yes"}), + ) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "`highlight` should be a boolean" in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q", "highlight": 1}), + ) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "`highlight` should be a boolean" in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q"}), + ) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (False, None)) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert "Dataset not found!" in res["message"] + + feature_calls = {"cross": None, "keyword": None, "retrieval_question": None} + + async def _cross_languages(_tenant_id, _dialog, question, langs): + feature_calls["cross"] = tuple(langs) + return f"{question}-xl" + + async def _keyword_extraction(_chat_mdl, question): + feature_calls["keyword"] = question + return "-kw" + + class _FeatureRetriever: + async def retrieval(self, question, *_args, **_kwargs): + feature_calls["retrieval_question"] = question + return { + "chunks": [ + { + "chunk_id": "c1", + "content_with_weight": "content", + "doc_id": "doc-1", + "kb_id": "ds-1", + "vector": [1, 2], + } + ], + "total": 1, + } + + async def retrieval_by_toc(self, question, chunks, tenant_ids, _chat_mdl, size): + assert question == "q-xl-kw" + assert chunks and tenant_ids + assert size == 30 + return [ + { + "chunk_id": "toc-1", + "content_with_weight": "toc content", + "doc_id": "doc-toc", + "kb_id": "ds-1", + } + ] + + def retrieval_by_children(self, chunks, _tenant_ids): + return chunks + [ + { + "chunk_id": "child-1", + "content_with_weight": "child content", + "doc_id": "doc-child", + "kb_id": "ds-1", + } + ] + + class _FeatureKgRetriever: + async def retrieval(self, *_args, **_kwargs): + return { + "chunk_id": "kg-1", + "content_with_weight": "kg content", + "doc_id": "doc-kg", + "kb_id": "ds-1", + } + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue( + { + "dataset_ids": ["ds-1"], + "question": "q", + "rerank_id": "rerank-1", + "cross_languages": ["fr"], + "keyword": True, + "toc_enhance": True, + "use_kg": True, + } + ), + ) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", embd_id="m1"))) + monkeypatch.setattr(module, "cross_languages", _cross_languages) + monkeypatch.setattr(module, "keyword_extraction", _keyword_extraction) + monkeypatch.setattr(module.settings, "retriever", _FeatureRetriever()) + monkeypatch.setattr(module.settings, "kg_retriever", _FeatureKgRetriever()) + monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: {}) + monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: SimpleNamespace()) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert res["code"] == 0 + assert feature_calls["cross"] == ("fr",) + assert feature_calls["keyword"] == "q-xl" + assert feature_calls["retrieval_question"] == "q-xl-kw" + assert res["data"]["chunks"][0]["id"] == "kg-1" + assert res["data"]["chunks"][0]["content"] == "kg content" + assert any(chunk["id"] == "toc-1" for chunk in res["data"]["chunks"]) + assert any(chunk["id"] == "child-1" for chunk in res["data"]["chunks"]) + + class _NotFoundRetriever: + async def retrieval(self, *_args, **_kwargs): + raise Exception("boom not_found boom") + + def retrieval_by_children(self, chunks, *_args, **_kwargs): + return chunks + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q"}), + ) + monkeypatch.setattr(module.settings, "retriever", _NotFoundRetriever()) + res = _run(module.retrieval_test.__wrapped__("tenant-1")) + assert res["code"] == module.RetCode.DATA_ERROR + assert "No chunk found! Check the chunk status please!" in res["message"] diff --git a/test/testcases/test_http_api/test_session_management/test_agent_sessions.py b/test/testcases/test_http_api/test_session_management/test_agent_sessions.py index 6f1d65fa5e..cfcc1807be 100644 --- a/test/testcases/test_http_api/test_session_management/test_agent_sessions.py +++ b/test/testcases/test_http_api/test_session_management/test_agent_sessions.py @@ -14,6 +14,7 @@ # limitations under the License. # import pytest +import requests from common import ( create_agent, create_agent_session, @@ -22,6 +23,7 @@ from common import ( list_agent_sessions, list_agents, ) +from configs import HOST_ADDRESS, VERSION AGENT_TITLE = "test_agent_http" MINIMAL_DSL = { @@ -87,3 +89,33 @@ class TestAgentSessions: res = delete_agent_sessions(HttpApiAuth, agent_id, {"ids": [session_id]}) assert res["code"] == 0, res + + @pytest.mark.p2 + def test_agent_crud_validation_contract(self, HttpApiAuth, agent_id): + res = list_agents(HttpApiAuth, {"id": "missing-agent-id", "title": "missing-agent-title"}) + assert res["code"] == 102, res + assert "doesn't exist" in res["message"], res + + res = list_agents(HttpApiAuth, {"title": AGENT_TITLE, "desc": "true", "page_size": 1}) + assert res["code"] == 0, res + + res = create_agent(HttpApiAuth, {"title": "missing-dsl-agent"}) + assert res["code"] == 101, res + assert "No DSL data in request" in res["message"], res + + res = create_agent(HttpApiAuth, {"dsl": MINIMAL_DSL}) + assert res["code"] == 101, res + assert "No title in request" in res["message"], res + + res = create_agent(HttpApiAuth, {"title": AGENT_TITLE, "dsl": MINIMAL_DSL}) + assert res["code"] == 102, res + assert "already exists" in res["message"], res + + update_url = f"{HOST_ADDRESS}/api/{VERSION}/agents/invalid-agent-id" + res = requests.put(update_url, auth=HttpApiAuth, json={"title": "updated", "dsl": MINIMAL_DSL}).json() + assert res["code"] == 103, res + assert "Only owner of canvas authorized" in res["message"], res + + res = delete_agent(HttpApiAuth, "invalid-agent-id") + assert res["code"] == 103, res + assert "Only owner of canvas authorized" in res["message"], res diff --git a/test/testcases/test_http_api/test_session_management/test_chat_completions_openai.py b/test/testcases/test_http_api/test_session_management/test_chat_completions_openai.py index e126119ad1..ffaa3ee451 100644 --- a/test/testcases/test_http_api/test_session_management/test_chat_completions_openai.py +++ b/test/testcases/test_http_api/test_session_management/test_chat_completions_openai.py @@ -130,3 +130,80 @@ class TestChatCompletionsOpenAI: ) # Should return an error (format may vary based on implementation) assert "error" in res or res.get("code") != 0, f"Should return error for invalid chat: {res}" + + @pytest.mark.p2 + @pytest.mark.parametrize( + "payload, requires_valid_chat, expected_message", + [ + ( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "extra_body": "invalid_extra_body", + }, + False, + "extra_body must be an object.", + ), + ( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "extra_body": {"reference_metadata": "invalid_reference_metadata"}, + }, + False, + "reference_metadata must be an object.", + ), + ( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "extra_body": {"reference_metadata": {"fields": "author"}}, + }, + False, + "reference_metadata.fields must be an array.", + ), + ( + { + "model": "model", + "messages": [], + }, + False, + "You have to provide messages.", + ), + ( + { + "model": "model", + "messages": [{"role": "assistant", "content": "hello"}], + }, + False, + "The last content of this conversation is not from user.", + ), + ( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "extra_body": {"metadata_condition": "invalid"}, + }, + True, + "metadata_condition must be an object.", + ), + ], + ) + def test_openai_chat_completion_request_validation( + self, + HttpApiAuth, + request, + payload, + requires_valid_chat, + expected_message, + ): + chat_id = "invalid_chat_id" + if requires_valid_chat: + res = create_chat_assistant(HttpApiAuth, {"name": "openai_validation_case", "dataset_ids": []}) + assert res["code"] == 0, res + chat_id = res["data"]["id"] + request.addfinalizer(lambda: delete_chat_assistants(HttpApiAuth)) + + res = chat_completions_openai(HttpApiAuth, chat_id, payload) + assert res.get("code") != 0, res + assert expected_message in res.get("message", ""), res diff --git a/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py b/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py new file mode 100644 index 0000000000..7660ae6c36 --- /dev/null +++ b/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py @@ -0,0 +1,1127 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import inspect +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _Args(dict): + def get(self, key, default=None, type=None): + value = super().get(key, default) + if value is None or type is None: + return value + try: + return type(value) + except (TypeError, ValueError): + return default + + +class _StubHeaders: + def __init__(self): + self._items = [] + + def add_header(self, key, value): + self._items.append((key, value)) + + def get(self, key, default=None): + for existing_key, value in reversed(self._items): + if existing_key == key: + return value + return default + + +class _StubResponse: + def __init__(self, body, mimetype=None, content_type=None): + self.body = body + self.mimetype = mimetype + self.content_type = content_type + self.headers = _StubHeaders() + + +def _run(coro): + return asyncio.run(coro) + + +async def _collect_stream(body): + items = [] + if hasattr(body, "__aiter__"): + async for item in body: + if isinstance(item, bytes): + item = item.decode("utf-8") + items.append(item) + else: + for item in body: + if isinstance(item, bytes): + item = item.decode("utf-8") + items.append(item) + return items + + +def _load_session_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + deepdoc_pkg = ModuleType("deepdoc") + deepdoc_parser_pkg = ModuleType("deepdoc.parser") + deepdoc_parser_pkg.__path__ = [] + + class _StubPdfParser: + pass + + class _StubExcelParser: + pass + + class _StubDocxParser: + pass + + deepdoc_parser_pkg.PdfParser = _StubPdfParser + deepdoc_parser_pkg.ExcelParser = _StubExcelParser + deepdoc_parser_pkg.DocxParser = _StubDocxParser + deepdoc_pkg.parser = deepdoc_parser_pkg + monkeypatch.setitem(sys.modules, "deepdoc", deepdoc_pkg) + monkeypatch.setitem(sys.modules, "deepdoc.parser", deepdoc_parser_pkg) + + deepdoc_excel_module = ModuleType("deepdoc.parser.excel_parser") + deepdoc_excel_module.RAGFlowExcelParser = _StubExcelParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.excel_parser", deepdoc_excel_module) + + deepdoc_mineru_module = ModuleType("deepdoc.parser.mineru_parser") + + class _StubMinerUParser: + pass + + deepdoc_mineru_module.MinerUParser = _StubMinerUParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.mineru_parser", deepdoc_mineru_module) + + deepdoc_paddle_module = ModuleType("deepdoc.parser.paddleocr_parser") + + class _StubPaddleOCRParser: + pass + + deepdoc_paddle_module.PaddleOCRParser = _StubPaddleOCRParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.paddleocr_parser", deepdoc_paddle_module) + + deepdoc_parser_utils = ModuleType("deepdoc.parser.utils") + deepdoc_parser_utils.get_text = lambda *_args, **_kwargs: "" + monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) + monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + + agent_pkg = ModuleType("agent") + agent_pkg.__path__ = [] + agent_canvas_mod = ModuleType("agent.canvas") + + class _StubCanvas: + def __init__(self, *_args, **_kwargs): + self._dsl = "{}" + + def reset(self): + return None + + def get_prologue(self): + return "stub prologue" + + def get_component_input_form(self, _name): + return {} + + def get_mode(self): + return "chat" + + def __str__(self): + return self._dsl + + agent_canvas_mod.Canvas = _StubCanvas + agent_pkg.canvas = agent_canvas_mod + monkeypatch.setitem(sys.modules, "agent", agent_pkg) + monkeypatch.setitem(sys.modules, "agent.canvas", agent_canvas_mod) + + module_path = repo_root / "api" / "apps" / "sdk" / "session.py" + spec = importlib.util.spec_from_file_location("test_session_sdk_routes_unit_module", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, "test_session_sdk_routes_unit_module", module) + spec.loader.exec_module(module) + return module + + +@pytest.mark.p2 +def test_create_and_update_guard_matrix(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"name": "session"})) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.create)("tenant-1", "chat-1")) + assert res["message"] == "You do not own the assistant." + + dia = SimpleNamespace(prompt_config={"prologue": "hello"}) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [dia]) + monkeypatch.setattr(module.ConversationService, "save", lambda **_kwargs: None) + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (False, None)) + res = _run(inspect.unwrap(module.create)("tenant-1", "chat-1")) + assert "Fail to create a session" in res["message"] + + monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args())) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (False, None)) + res = _run(inspect.unwrap(module.create_agent_session)("tenant-1", "agent-1")) + assert res["message"] == "Agent not found." + + canvas = SimpleNamespace(dsl="{}", id="agent-1") + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, canvas)) + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.create_agent_session)("tenant-1", "agent-1")) + assert res["message"] == "You cannot access the agent." + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({})) + monkeypatch.setattr(module.ConversationService, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.update)("tenant-1", "chat-1", "session-1")) + assert res["message"] == "Session does not exist" + + monkeypatch.setattr(module.ConversationService, "query", lambda **_kwargs: [SimpleNamespace(id="session-1")]) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.update)("tenant-1", "chat-1", "session-1")) + assert res["message"] == "You do not own the session" + + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="chat-1")]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"message": []})) + res = _run(inspect.unwrap(module.update)("tenant-1", "chat-1", "session-1")) + assert "`message` can not be change" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"reference": []})) + res = _run(inspect.unwrap(module.update)("tenant-1", "chat-1", "session-1")) + assert "`reference` can not be change" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"name": ""})) + res = _run(inspect.unwrap(module.update)("tenant-1", "chat-1", "session-1")) + assert "`name` can not be empty" in res["message"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"name": "renamed"})) + monkeypatch.setattr(module.ConversationService, "update_by_id", lambda *_args, **_kwargs: False) + res = _run(inspect.unwrap(module.update)("tenant-1", "chat-1", "session-1")) + assert res["message"] == "Session updates error" + + +@pytest.mark.p2 +def test_chat_completion_metadata_and_stream_paths(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "Response", _StubResponse) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(kb_ids=["kb-1"])]) + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda _kb_ids: [{"id": "doc-1"}]) + monkeypatch.setattr(module, "convert_conditions", lambda cond: cond.get("conditions", [])) + monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: []) + + captured_requests = [] + + async def fake_rag_completion(_tenant_id, _chat_id, **req): + captured_requests.append(req) + yield {"answer": "ok"} + + monkeypatch.setattr(module, "rag_completion", fake_rag_completion) + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(None)) + resp = _run(inspect.unwrap(module.chat_completion)("tenant-1", "chat-1")) + assert isinstance(resp, _StubResponse) + assert resp.headers.get("Content-Type") == "text/event-stream; charset=utf-8" + _run(_collect_stream(resp.body)) + assert captured_requests[-1].get("question") == "" + + req_with_conditions = { + "question": "hello", + "session_id": "session-1", + "metadata_condition": {"logic": "and", "conditions": [{"name": "author", "value": "bob"}]}, + "stream": True, + } + monkeypatch.setattr(module.ConversationService, "query", lambda **_kwargs: [SimpleNamespace(id="session-1")]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(req_with_conditions)) + resp = _run(inspect.unwrap(module.chat_completion)("tenant-1", "chat-1")) + _run(_collect_stream(resp.body)) + assert captured_requests[-1].get("doc_ids") == "-999" + + req_without_conditions = { + "question": "hello", + "session_id": "session-1", + "metadata_condition": {"logic": "and", "conditions": []}, + "stream": True, + "doc_ids": "legacy", + } + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(req_without_conditions)) + resp = _run(inspect.unwrap(module.chat_completion)("tenant-1", "chat-1")) + _run(_collect_stream(resp.body)) + assert "doc_ids" not in captured_requests[-1] + + +@pytest.mark.p2 +def test_openai_chat_validation_matrix_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "num_tokens_from_string", lambda _text: 1) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(kb_ids=["kb-1"])]) + + cases = [ + ( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "extra_body": "bad", + }, + "extra_body must be an object.", + ), + ( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "extra_body": {"reference_metadata": "bad"}, + }, + "reference_metadata must be an object.", + ), + ( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "extra_body": {"reference_metadata": {"fields": "bad"}}, + }, + "reference_metadata.fields must be an array.", + ), + ({"model": "model", "messages": []}, "You have to provide messages."), + ( + {"model": "model", "messages": [{"role": "assistant", "content": "hello"}]}, + "The last content of this conversation is not from user.", + ), + ( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "extra_body": {"metadata_condition": "bad"}, + }, + "metadata_condition must be an object.", + ), + ] + + for payload, expected in cases: + monkeypatch.setattr(module, "get_request_json", lambda p=payload: _AwaitableValue(p)) + res = _run(inspect.unwrap(module.chat_completion_openai_like)("tenant-1", "chat-1")) + assert expected in res["message"] + + +@pytest.mark.p2 +def test_openai_stream_generator_branches_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "Response", _StubResponse) + monkeypatch.setattr(module, "num_tokens_from_string", lambda text: len(text or "")) + monkeypatch.setattr(module, "convert_conditions", lambda cond: cond.get("conditions", [])) + monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: []) + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda _kb_ids: [{"id": "doc-1"}]) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(kb_ids=["kb-1"])]) + monkeypatch.setattr(module, "_build_reference_chunks", lambda *_args, **_kwargs: [{"id": "ref-1"}]) + + async def fake_async_chat(_dia, _msg, _stream, **_kwargs): + yield {"start_to_think": True} + yield {"answer": "R"} + yield {"end_to_think": True} + yield {"answer": ""} + yield {"answer": "C"} + yield {"final": True, "answer": "DONE", "reference": {"chunks": []}} + raise RuntimeError("boom") + + monkeypatch.setattr(module, "async_chat", fake_async_chat) + + payload = { + "model": "model", + "stream": True, + "messages": [ + {"role": "system", "content": "sys"}, + {"role": "assistant", "content": "preface"}, + {"role": "user", "content": "hello"}, + ], + "extra_body": { + "reference": True, + "reference_metadata": {"include": True, "fields": ["author"]}, + "metadata_condition": {"logic": "and", "conditions": [{"name": "author", "value": "bob"}]}, + }, + } + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(payload)) + + resp = _run(inspect.unwrap(module.chat_completion_openai_like)("tenant-1", "chat-1")) + assert isinstance(resp, _StubResponse) + assert resp.headers.get("Content-Type") == "text/event-stream; charset=utf-8" + + chunks = _run(_collect_stream(resp.body)) + assert any("reasoning_content" in chunk for chunk in chunks) + assert any("**ERROR**: boom" in chunk for chunk in chunks) + assert any('"usage"' in chunk for chunk in chunks) + assert any('"reference"' in chunk for chunk in chunks) + assert chunks[-1].strip() == "data:[DONE]" + + +@pytest.mark.p2 +def test_openai_nonstream_branch_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "jsonify", lambda payload: payload) + monkeypatch.setattr(module, "num_tokens_from_string", lambda text: len(text or "")) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(kb_ids=[])]) + + async def fake_async_chat(_dia, _msg, _stream, **_kwargs): + yield {"answer": "world", "reference": {}} + + monkeypatch.setattr(module, "async_chat", fake_async_chat) + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue( + { + "model": "model", + "messages": [{"role": "user", "content": "hello"}], + "stream": False, + } + ), + ) + + res = _run(inspect.unwrap(module.chat_completion_openai_like)("tenant-1", "chat-1")) + assert res["choices"][0]["message"]["content"] == "world" + + +@pytest.mark.p2 +def test_agents_openai_compatibility_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "Response", _StubResponse) + monkeypatch.setattr(module, "jsonify", lambda payload: payload) + monkeypatch.setattr(module, "num_tokens_from_string", lambda text: len(text or "")) + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"model": "model", "messages": []})) + res = _run(inspect.unwrap(module.agents_completion_openai_compatibility)("tenant-1", "agent-1")) + assert "at least one message" in res["message"] + + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"model": "model", "messages": [{"role": "user", "content": "hello"}]}), + ) + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.agents_completion_openai_compatibility)("tenant-1", "agent-1")) + assert "don't own the agent" in res["message"] + + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: [SimpleNamespace(id="agent-1")]) + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"model": "model", "messages": [{"role": "system", "content": "system only"}]}), + ) + res = _run(inspect.unwrap(module.agents_completion_openai_compatibility)("tenant-1", "agent-1")) + assert "No valid messages found" in json.dumps(res) + + captured_calls = [] + + async def _completion_openai_stream(*args, **kwargs): + captured_calls.append((args, kwargs)) + yield "data:stream" + + monkeypatch.setattr(module, "completion_openai", _completion_openai_stream) + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue( + { + "model": "model", + "messages": [ + {"role": "assistant", "content": "preface"}, + {"role": "user", "content": "latest question"}, + ], + "stream": True, + "metadata": {"id": "meta-session"}, + } + ), + ) + resp = _run(inspect.unwrap(module.agents_completion_openai_compatibility)("tenant-1", "agent-1")) + assert isinstance(resp, _StubResponse) + assert resp.headers.get("Content-Type") == "text/event-stream; charset=utf-8" + _run(_collect_stream(resp.body)) + assert captured_calls[-1][0][2] == "latest question" + + async def _completion_openai_nonstream(*args, **kwargs): + captured_calls.append((args, kwargs)) + yield {"id": "non-stream"} + + monkeypatch.setattr(module, "completion_openai", _completion_openai_nonstream) + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue( + { + "model": "model", + "messages": [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "middle"}, + {"role": "user", "content": "final user"}, + ], + "stream": False, + "session_id": "session-1", + "temperature": 0.5, + } + ), + ) + res = _run(inspect.unwrap(module.agents_completion_openai_compatibility)("tenant-1", "agent-1")) + assert res["id"] == "non-stream" + assert captured_calls[-1][0][2] == "final user" + assert captured_calls[-1][1]["stream"] is False + assert captured_calls[-1][1]["session_id"] == "session-1" + + +@pytest.mark.p2 +def test_agent_completions_stream_and_nonstream_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "Response", _StubResponse) + + async def _agent_stream(*_args, **_kwargs): + yield "data:not-json" + yield "data:" + json.dumps({"event": "node_finished", "data": {"component_id": "c1"}}) + yield "data:" + json.dumps({"event": "other", "data": {}}) + yield "data:" + json.dumps({"event": "message", "data": {"content": "hello"}}) + + monkeypatch.setattr(module, "agent_completion", _agent_stream) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"stream": True, "return_trace": True})) + + resp = _run(inspect.unwrap(module.agent_completions)("tenant-1", "agent-1")) + chunks = _run(_collect_stream(resp.body)) + assert resp.headers.get("Content-Type") == "text/event-stream; charset=utf-8" + assert any('"trace"' in chunk for chunk in chunks) + assert any("hello" in chunk for chunk in chunks) + assert chunks[-1].strip() == "data:[DONE]" + + async def _agent_nonstream(*_args, **_kwargs): + yield "data:" + json.dumps({"event": "message", "data": {"content": "A", "reference": {"doc": "r"}}}) + yield "data:" + json.dumps({"event": "node_finished", "data": {"component_id": "c2"}}) + + monkeypatch.setattr(module, "agent_completion", _agent_nonstream) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"stream": False, "return_trace": True})) + res = _run(inspect.unwrap(module.agent_completions)("tenant-1", "agent-1")) + assert res["data"]["data"]["content"] == "A" + assert res["data"]["data"]["reference"] == {"doc": "r"} + assert res["data"]["data"]["trace"][0]["component_id"] == "c2" + + async def _agent_nonstream_broken(*_args, **_kwargs): + yield "data:{" + + monkeypatch.setattr(module, "agent_completion", _agent_nonstream_broken) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"stream": False, "return_trace": False})) + res = _run(inspect.unwrap(module.agent_completions)("tenant-1", "agent-1")) + assert res["data"].startswith("**ERROR**") + + +@pytest.mark.p2 +def test_list_session_projection_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args({}))) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="chat-1")]) + + convs = [ + { + "id": "session-1", + "dialog_id": "chat-1", + "message": [{"role": "assistant", "content": "hello", "prompt": "internal"}], + "reference": [ + { + "chunks": [ + { + "chunk_id": "chunk-1", + "content_with_weight": "weighted", + "doc_id": "doc-1", + "docnm_kwd": "doc-name", + "kb_id": "kb-1", + "image_id": "img-1", + "positions": [1, 2], + } + ] + } + ], + } + ] + monkeypatch.setattr(module.ConversationService, "get_list", lambda *_args, **_kwargs: convs) + + res = _run(inspect.unwrap(module.list_session)("tenant-1", "chat-1")) + assert res["data"][0]["chat_id"] == "chat-1" + assert "reference" not in res["data"][0] + assert "prompt" not in res["data"][0]["messages"][0] + assert res["data"][0]["messages"][0]["reference"][0]["positions"] == [1, 2] + + +@pytest.mark.p2 +def test_list_agent_session_projection_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args({}))) + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: [SimpleNamespace(id="agent-1")]) + + conv_non_list_reference = { + "id": "session-1", + "dialog_id": "agent-1", + "message": [{"role": "assistant", "content": "hello", "prompt": "internal"}], + "reference": {"unexpected": "shape"}, + } + monkeypatch.setattr(module.API4ConversationService, "get_list", lambda *_args, **_kwargs: (1, [conv_non_list_reference])) + res = _run(inspect.unwrap(module.list_agent_session)("tenant-1", "agent-1")) + assert res["data"][0]["agent_id"] == "agent-1" + assert "prompt" not in res["data"][0]["messages"][0] + + conv_with_chunks = { + "id": "session-2", + "dialog_id": "agent-1", + "message": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer", "prompt": "internal"}, + ], + "reference": [ + { + "chunks": [ + "not-a-dict", + { + "chunk_id": "chunk-2", + "content_with_weight": "weighted", + "doc_id": "doc-2", + "docnm_kwd": "doc-name-2", + "kb_id": "kb-2", + "image_id": "img-2", + "positions": [9], + }, + ] + } + ], + } + monkeypatch.setattr(module.API4ConversationService, "get_list", lambda *_args, **_kwargs: (1, [conv_with_chunks])) + res = _run(inspect.unwrap(module.list_agent_session)("tenant-1", "agent-1")) + projected_chunk = res["data"][0]["messages"][1]["reference"][0] + assert projected_chunk["image_id"] == "img-2" + assert projected_chunk["positions"] == [9] + + +@pytest.mark.p2 +def test_delete_routes_partial_duplicate_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="chat-1")]) + monkeypatch.setattr(module.ConversationService, "delete_by_id", lambda *_args, **_kwargs: True) + + def _conversation_query(**kwargs): + if "id" not in kwargs: + return [SimpleNamespace(id="seed")] + if kwargs["id"] == "ok": + return [SimpleNamespace(id="ok")] + return [] + + monkeypatch.setattr(module.ConversationService, "query", _conversation_query) + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"ids": ["ok", "bad"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + res = _run(inspect.unwrap(module.delete)("tenant-1", "chat-1")) + assert res["code"] == 0 + assert res["data"]["success_count"] == 1 + assert res["data"]["errors"] == ["The chat doesn't own the session bad"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"ids": ["bad"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + res = _run(inspect.unwrap(module.delete)("tenant-1", "chat-1")) + assert res["message"] == "The chat doesn't own the session bad" + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"ids": ["ok", "ok"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (["ok"], ["Duplicate session ids: ok"])) + res = _run(inspect.unwrap(module.delete)("tenant-1", "chat-1")) + assert res["code"] == 0 + assert res["data"]["success_count"] == 1 + assert res["data"]["errors"] == ["Duplicate session ids: ok"] + + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: [SimpleNamespace(id="agent-1")]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"ids": ["session-1"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + + def _agent_query(**kwargs): + if "id" not in kwargs: + return [SimpleNamespace(id="session-1")] + if kwargs["id"] == "session-1": + return [SimpleNamespace(id="session-1")] + return [] + + monkeypatch.setattr(module.API4ConversationService, "query", _agent_query) + monkeypatch.setattr(module.API4ConversationService, "delete_by_id", lambda *_args, **_kwargs: True) + res = _run(inspect.unwrap(module.delete_agent_session)("tenant-1", "agent-1")) + assert res["code"] == 0 + + +@pytest.mark.p2 +def test_delete_agent_session_error_matrix_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: [SimpleNamespace(id="agent-1")]) + monkeypatch.setattr(module.API4ConversationService, "delete_by_id", lambda *_args, **_kwargs: True) + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"ids": ["ok", "missing"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + + def _query_partial(**kwargs): + if "id" not in kwargs: + return [SimpleNamespace(id="ok"), SimpleNamespace(id="missing")] + if kwargs["id"] == "ok": + return [SimpleNamespace(id="ok")] + return [] + + monkeypatch.setattr(module.API4ConversationService, "query", _query_partial) + res = _run(inspect.unwrap(module.delete_agent_session)("tenant-1", "agent-1")) + assert res["data"]["success_count"] == 1 + assert res["data"]["errors"] == ["The agent doesn't own the session missing"] + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"ids": ["missing"]})) + + def _query_all_failed(**kwargs): + if "id" not in kwargs: + return [SimpleNamespace(id="missing")] + return [] + + monkeypatch.setattr(module.API4ConversationService, "query", _query_all_failed) + res = _run(inspect.unwrap(module.delete_agent_session)("tenant-1", "agent-1")) + assert res["message"] == "The agent doesn't own the session missing" + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"ids": ["ok", "ok"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (["ok"], ["Duplicate session ids: ok"])) + + def _query_duplicate(**kwargs): + if "id" not in kwargs: + return [SimpleNamespace(id="ok")] + if kwargs["id"] == "ok": + return [SimpleNamespace(id="ok")] + return [] + + monkeypatch.setattr(module.API4ConversationService, "query", _query_duplicate) + res = _run(inspect.unwrap(module.delete_agent_session)("tenant-1", "agent-1")) + assert res["data"]["success_count"] == 1 + assert res["data"]["errors"] == ["Duplicate session ids: ok"] + + +@pytest.mark.p2 +def test_sessions_ask_route_validation_and_stream_unit(monkeypatch): + module = _load_session_module(monkeypatch) + monkeypatch.setattr(module, "Response", _StubResponse) + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"dataset_ids": ["kb-1"]})) + res = _run(inspect.unwrap(module.ask_about)("tenant-1")) + assert res["message"] == "`question` is required." + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"question": "q"})) + res = _run(inspect.unwrap(module.ask_about)("tenant-1")) + assert res["message"] == "`dataset_ids` is required." + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"question": "q", "dataset_ids": "kb-1"})) + res = _run(inspect.unwrap(module.ask_about)("tenant-1")) + assert res["message"] == "`dataset_ids` should be a list." + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"question": "q", "dataset_ids": ["kb-1"]})) + res = _run(inspect.unwrap(module.ask_about)("tenant-1")) + assert res["message"] == "You don't own the dataset kb-1." + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [SimpleNamespace(chunk_num=0)]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"question": "q", "dataset_ids": ["kb-1"]})) + res = _run(inspect.unwrap(module.ask_about)("tenant-1")) + assert res["message"] == "The dataset kb-1 doesn't own parsed file" + + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [SimpleNamespace(chunk_num=1)]) + captured = {} + + async def _streaming_async_ask(question, kb_ids, uid): + captured["question"] = question + captured["kb_ids"] = kb_ids + captured["uid"] = uid + yield {"answer": "first"} + raise RuntimeError("ask stream boom") + + monkeypatch.setattr(module, "async_ask", _streaming_async_ask) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"question": "q", "dataset_ids": ["kb-1"]})) + resp = _run(inspect.unwrap(module.ask_about)("tenant-1")) + assert isinstance(resp, _StubResponse) + assert resp.headers.get("Content-Type") == "text/event-stream; charset=utf-8" + chunks = _run(_collect_stream(resp.body)) + assert any('"answer": "first"' in chunk for chunk in chunks) + assert any('"code": 500' in chunk and "**ERROR**: ask stream boom" in chunk for chunk in chunks) + assert '"data": true' in chunks[-1].lower() + assert captured == {"question": "q", "kb_ids": ["kb-1"], "uid": "tenant-1"} + + +@pytest.mark.p2 +def test_sessions_related_questions_prompt_build_unit(monkeypatch): + module = _load_session_module(monkeypatch) + + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({})) + res = _run(inspect.unwrap(module.related_questions)("tenant-1")) + assert res["message"] == "`question` is required." + + captured = {} + + class _FakeLLMBundle: + def __init__(self, *args, **kwargs): + captured["bundle_args"] = args + captured["bundle_kwargs"] = kwargs + + async def async_chat(self, prompt, messages, options): + captured["prompt"] = prompt + captured["messages"] = messages + captured["options"] = options + return "1. First related\n2. Second related\nplain text" + + monkeypatch.setattr(module, "LLMBundle", _FakeLLMBundle) + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"question": "solar energy", "industry": "renewables"}), + ) + res = _run(inspect.unwrap(module.related_questions)("tenant-1")) + assert res["data"] == ["First related", "Second related"] + assert "Keep the term length between 2-4 words" in captured["prompt"] + assert "related terms can also help search engines" in captured["prompt"] + assert "Ensure all search terms are relevant to the industry: renewables." in captured["prompt"] + assert "Keywords: solar energy" in captured["messages"][0]["content"] + assert captured["options"] == {"temperature": 0.9} + + +@pytest.mark.p2 +def test_chatbot_routes_auth_stream_nonstream_unit(monkeypatch): + module = _load_session_module(monkeypatch) + monkeypatch.setattr(module, "Response", _StubResponse) + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"})) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({})) + res = _run(inspect.unwrap(module.chatbot_completions)("dialog-1")) + assert res["message"] == "Authorization is not valid!" + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer bad"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.chatbot_completions)("dialog-1")) + assert "API key is invalid" in res["message"] + + stream_calls = [] + + async def _iframe_stream(dialog_id, **req): + stream_calls.append((dialog_id, dict(req))) + yield "data:stream-chunk" + + monkeypatch.setattr(module, "iframe_completion", _iframe_stream) + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer ok"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"stream": True})) + resp = _run(inspect.unwrap(module.chatbot_completions)("dialog-1")) + assert isinstance(resp, _StubResponse) + assert resp.headers.get("Content-Type") == "text/event-stream; charset=utf-8" + _run(_collect_stream(resp.body)) + assert stream_calls[-1][0] == "dialog-1" + assert stream_calls[-1][1]["quote"] is False + + async def _iframe_nonstream(_dialog_id, **_req): + yield {"answer": "non-stream"} + + monkeypatch.setattr(module, "iframe_completion", _iframe_nonstream) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"stream": False, "quote": True})) + res = _run(inspect.unwrap(module.chatbot_completions)("dialog-1")) + assert res["data"]["answer"] == "non-stream" + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"})) + res = _run(inspect.unwrap(module.chatbots_inputs)("dialog-1")) + assert res["message"] == "Authorization is not valid!" + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer invalid"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.chatbots_inputs)("dialog-1")) + assert "API key is invalid" in res["message"] + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer ok"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _dialog_id: (False, None)) + res = _run(inspect.unwrap(module.chatbots_inputs)("dialog-404")) + assert res["message"] == "Can't find dialog by ID: dialog-404" + + +@pytest.mark.p2 +def test_agentbot_routes_auth_stream_nonstream_unit(monkeypatch): + module = _load_session_module(monkeypatch) + monkeypatch.setattr(module, "Response", _StubResponse) + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"})) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({})) + res = _run(inspect.unwrap(module.agent_bot_completions)("agent-1")) + assert res["message"] == "Authorization is not valid!" + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer bad"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.agent_bot_completions)("agent-1")) + assert "API key is invalid" in res["message"] + + async def _agent_stream(*_args, **_kwargs): + yield "data:agent-stream" + + monkeypatch.setattr(module, "agent_completion", _agent_stream) + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer ok"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"stream": True})) + resp = _run(inspect.unwrap(module.agent_bot_completions)("agent-1")) + assert isinstance(resp, _StubResponse) + assert resp.headers.get("Content-Type") == "text/event-stream; charset=utf-8" + _run(_collect_stream(resp.body)) + + async def _agent_nonstream(*_args, **_kwargs): + yield {"answer": "agent-non-stream"} + + monkeypatch.setattr(module, "agent_completion", _agent_nonstream) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"stream": False})) + res = _run(inspect.unwrap(module.agent_bot_completions)("agent-1")) + assert res["data"]["answer"] == "agent-non-stream" + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"})) + res = _run(inspect.unwrap(module.begin_inputs)("agent-1")) + assert res["message"] == "Authorization is not valid!" + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer bad"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.begin_inputs)("agent-1")) + assert "API key is invalid" in res["message"] + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer ok"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _agent_id: (False, None)) + res = _run(inspect.unwrap(module.begin_inputs)("agent-404")) + assert res["message"] == "Can't find agent by ID: agent-404" + + +@pytest.mark.p2 +def test_searchbots_ask_embedded_auth_and_stream_unit(monkeypatch): + module = _load_session_module(monkeypatch) + monkeypatch.setattr(module, "Response", _StubResponse) + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"})) + res = _run(inspect.unwrap(module.ask_about_embedded)()) + assert res["message"] == "Authorization is not valid!" + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer bad"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.ask_about_embedded)()) + assert "API key is invalid" in res["message"] + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer ok"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"question": "embedded q", "kb_ids": ["kb-1"], "search_id": "search-1"}), + ) + monkeypatch.setattr(module.SearchService, "get_detail", lambda _search_id: {"search_config": {"mode": "test"}}) + captured = {} + + async def _embedded_async_ask(question, kb_ids, uid, search_config=None): + captured["question"] = question + captured["kb_ids"] = kb_ids + captured["uid"] = uid + captured["search_config"] = search_config + yield {"answer": "embedded-answer"} + raise RuntimeError("embedded stream boom") + + monkeypatch.setattr(module, "async_ask", _embedded_async_ask) + resp = _run(inspect.unwrap(module.ask_about_embedded)()) + assert isinstance(resp, _StubResponse) + assert resp.headers.get("Content-Type") == "text/event-stream; charset=utf-8" + chunks = _run(_collect_stream(resp.body)) + assert any('"answer": "embedded-answer"' in chunk for chunk in chunks) + assert any('"code": 500' in chunk and "**ERROR**: embedded stream boom" in chunk for chunk in chunks) + assert '"data": true' in chunks[-1].lower() + assert captured["search_config"] == {"mode": "test"} + + +@pytest.mark.p2 +def test_searchbots_retrieval_test_embedded_matrix_unit(monkeypatch): + module = _load_session_module(monkeypatch) + handler = inspect.unwrap(module.retrieval_test_embedded) + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"})) + res = _run(handler()) + assert res["message"] == "Authorization is not valid!" + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer invalid"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) + res = _run(handler()) + assert "API key is invalid" in res["message"] + + monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer ok"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"kb_id": [], "question": "q"})) + res = _run(handler()) + assert res["message"] == "Please specify dataset firstly." + + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="")]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"kb_id": "kb-1", "question": "q"})) + res = _run(handler()) + assert res["message"] == "permission denined." + + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"kb_id": ["kb-no-access"], "question": "q"})) + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-a")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: []) + res = _run(handler()) + assert "Only owner of dataset authorized for this operation." in res["message"] + + llm_calls = [] + + def _fake_llm_bundle(tenant_id, llm_type, *args, **kwargs): + llm_calls.append((tenant_id, llm_type, args, kwargs)) + return SimpleNamespace(tenant_id=tenant_id, llm_type=llm_type, args=args, kwargs=kwargs) + + monkeypatch.setattr(module, "LLMBundle", _fake_llm_bundle) + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue({"kb_id": "kb-1", "question": "q", "meta_data_filter": {"method": "auto"}}), + ) + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda _kb_ids: [{"id": "doc-1"}]) + + async def _apply_filter(_meta_filter, _metas, _question, _chat_mdl, _local_doc_ids): + return ["doc-filtered"] + + monkeypatch.setattr(module, "apply_meta_data_filter", _apply_filter) + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-a")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [SimpleNamespace(id="kb-1")]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + res = _run(handler()) + assert res["message"] == "Knowledgebase not found!" + assert any(call[1] == module.LLMType.CHAT for call in llm_calls) + + llm_calls.clear() + retrieval_capture = {} + + async def _fake_retrieval( + question, + embd_mdl, + tenant_ids, + kb_ids, + page, + size, + similarity_threshold, + vector_similarity_weight, + top, + local_doc_ids, + rerank_mdl=None, + highlight=None, + rank_feature=None, + ): + retrieval_capture.update( + { + "question": question, + "embd_mdl": embd_mdl, + "tenant_ids": tenant_ids, + "kb_ids": kb_ids, + "page": page, + "size": size, + "similarity_threshold": similarity_threshold, + "vector_similarity_weight": vector_similarity_weight, + "top": top, + "local_doc_ids": local_doc_ids, + "rerank_mdl": rerank_mdl, + "highlight": highlight, + "rank_feature": rank_feature, + } + ) + return {"chunks": [{"id": "chunk-1", "vector": [0.1]}]} + + async def _translate(_tenant_id, _chat_id, question, _langs): + return question + "-translated" + + monkeypatch.setattr(module, "cross_languages", _translate) + monkeypatch.setattr(module, "label_question", lambda _question, _kbs: ["label-1"]) + monkeypatch.setattr(module.settings, "retriever", SimpleNamespace(retrieval=_fake_retrieval)) + monkeypatch.setattr( + module, + "get_request_json", + lambda: _AwaitableValue( + { + "kb_id": "kb-1", + "question": "translated-q", + "doc_ids": ["doc-seed"], + "cross_languages": ["es"], + "search_id": "search-1", + } + ), + ) + monkeypatch.setattr( + module.SearchService, + "get_detail", + lambda _search_id: { + "search_config": { + "meta_data_filter": {"method": "auto"}, + "chat_id": "chat-for-filter", + "similarity_threshold": 0.42, + "vector_similarity_weight": 0.8, + "top_k": 7, + "rerank_id": "reranker-model", + } + }, + ) + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda _kb_ids: [{"id": "doc-2"}]) + monkeypatch.setattr(module, "apply_meta_data_filter", _apply_filter) + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-a")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [SimpleNamespace(id="kb-1")]) + monkeypatch.setattr( + module.KnowledgebaseService, + "get_by_id", + lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-model")), + ) + res = _run(handler()) + assert res["code"] == 0 + assert res["data"]["labels"] == ["label-1"] + assert "vector" not in res["data"]["chunks"][0] + assert retrieval_capture["kb_ids"] == ["kb-1"] + assert retrieval_capture["tenant_ids"] == ["tenant-a"] + assert retrieval_capture["question"] == "translated-q-translated" + assert retrieval_capture["similarity_threshold"] == 0.42 + assert retrieval_capture["vector_similarity_weight"] == 0.8 + assert retrieval_capture["top"] == 7 + assert retrieval_capture["local_doc_ids"] == ["doc-filtered"] + assert retrieval_capture["rank_feature"] == ["label-1"] + assert retrieval_capture["rerank_mdl"] is not None + assert any(call[1] == module.LLMType.EMBEDDING.value and call[3].get("llm_name") == "embd-model" for call in llm_calls) diff --git a/test/testcases/test_sdk_api/test_agent_management/test_agent_crud_unit.py b/test/testcases/test_sdk_api/test_agent_management/test_agent_crud_unit.py new file mode 100644 index 0000000000..14fc03acf8 --- /dev/null +++ b/test/testcases/test_sdk_api/test_agent_management/test_agent_crud_unit.py @@ -0,0 +1,124 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import pytest +from ragflow_sdk import RAGFlow +from ragflow_sdk.modules.agent import Agent + + +class _DummyResponse: + def __init__(self, payload): + self._payload = payload + + def json(self): + return self._payload + + +@pytest.mark.p2 +def test_list_agents_success_and_error(monkeypatch): + client = RAGFlow("token", "http://localhost:9380") + captured = {} + + def _ok_get(path, params=None, json=None): + captured["path"] = path + captured["params"] = params + captured["json"] = json + return _DummyResponse({"code": 0, "data": [{"id": "agent-1", "title": "Agent One"}]}) + + monkeypatch.setattr(client, "get", _ok_get) + agents = client.list_agents(title="Agent One") + assert captured["path"] == "/agents" + assert captured["params"]["title"] == "Agent One" + assert isinstance(agents[0], Agent), str(agents) + assert agents[0].id == "agent-1", str(agents[0]) + assert agents[0].title == "Agent One", str(agents[0]) + + monkeypatch.setattr(client, "get", lambda *_args, **_kwargs: _DummyResponse({"code": 1, "message": "list boom"})) + with pytest.raises(Exception) as exception_info: + client.list_agents() + assert "list boom" in str(exception_info.value), str(exception_info.value) + + +@pytest.mark.p2 +def test_create_agent_payload_and_error(monkeypatch): + client = RAGFlow("token", "http://localhost:9380") + calls = [] + + def _ok_post(path, json=None, stream=False, files=None): + calls.append((path, json, stream, files)) + return _DummyResponse({"code": 0, "message": "ok"}) + + monkeypatch.setattr(client, "post", _ok_post) + client.create_agent("agent-title", {"graph": {}}, description=None) + assert calls[-1][0] == "/agents" + assert calls[-1][1] == {"title": "agent-title", "dsl": {"graph": {}}} + + client.create_agent("agent-title", {"graph": {}}, description="desc") + assert calls[-1][1] == {"title": "agent-title", "dsl": {"graph": {}}, "description": "desc"} + + monkeypatch.setattr(client, "post", lambda *_args, **_kwargs: _DummyResponse({"code": 1, "message": "create boom"})) + with pytest.raises(Exception) as exception_info: + client.create_agent("agent-title", {"graph": {}}) + assert "create boom" in str(exception_info.value), str(exception_info.value) + + +@pytest.mark.p2 +def test_update_agent_payload_matrix_and_error(monkeypatch): + client = RAGFlow("token", "http://localhost:9380") + calls = [] + + def _ok_put(path, json): + calls.append((path, json)) + return _DummyResponse({"code": 0, "message": "ok"}) + + monkeypatch.setattr(client, "put", _ok_put) + cases = [ + ({"title": "new-title"}, {"title": "new-title"}), + ({"description": "new-description"}, {"description": "new-description"}), + ({"dsl": {"nodes": []}}, {"dsl": {"nodes": []}}), + ( + {"title": "new-title", "description": "new-description", "dsl": {"nodes": []}}, + {"title": "new-title", "description": "new-description", "dsl": {"nodes": []}}, + ), + ] + for kwargs, expected_payload in cases: + client.update_agent("agent-1", **kwargs) + assert calls[-1][0] == "/agents/agent-1" + assert calls[-1][1] == expected_payload + + monkeypatch.setattr(client, "put", lambda *_args, **_kwargs: _DummyResponse({"code": 1, "message": "update boom"})) + with pytest.raises(Exception) as exception_info: + client.update_agent("agent-1", title="bad") + assert "update boom" in str(exception_info.value), str(exception_info.value) + + +@pytest.mark.p2 +def test_delete_agent_success_and_error(monkeypatch): + client = RAGFlow("token", "http://localhost:9380") + calls = [] + + def _ok_delete(path, json): + calls.append((path, json)) + return _DummyResponse({"code": 0, "message": "ok"}) + + monkeypatch.setattr(client, "delete", _ok_delete) + client.delete_agent("agent-1") + assert calls[-1] == ("/agents/agent-1", {}) + + monkeypatch.setattr(client, "delete", lambda *_args, **_kwargs: _DummyResponse({"code": 1, "message": "delete boom"})) + with pytest.raises(Exception) as exception_info: + client.delete_agent("agent-1") + assert "delete boom" in str(exception_info.value), str(exception_info.value) diff --git a/test/testcases/test_sdk_api/test_chat_assistant_management/test_delete_chat_assistants.py b/test/testcases/test_sdk_api/test_chat_assistant_management/test_delete_chat_assistants.py index 3d3be6f522..7f72033096 100644 --- a/test/testcases/test_sdk_api/test_chat_assistant_management/test_delete_chat_assistants.py +++ b/test/testcases/test_sdk_api/test_chat_assistant_management/test_delete_chat_assistants.py @@ -49,6 +49,17 @@ class TestChatAssistantsDelete: assistants = client.list_chats() assert len(assistants) == remaining + @pytest.mark.p2 + def test_delete_chats_nonzero_response_raises(self, client, monkeypatch): + class _DummyResponse: + def json(self): + return {"code": 1, "message": "boom"} + + monkeypatch.setattr(client, "delete", lambda *_args, **_kwargs: _DummyResponse()) + with pytest.raises(Exception) as exception_info: + client.delete_chats(ids=["chat-1"]) + assert "boom" in str(exception_info.value), str(exception_info.value) + @pytest.mark.parametrize( "payload", [ diff --git a/test/testcases/test_sdk_api/test_chat_assistant_management/test_update_chat_assistant.py b/test/testcases/test_sdk_api/test_chat_assistant_management/test_update_chat_assistant.py index df32561cc4..7652c266e7 100644 --- a/test/testcases/test_sdk_api/test_chat_assistant_management/test_update_chat_assistant.py +++ b/test/testcases/test_sdk_api/test_chat_assistant_management/test_update_chat_assistant.py @@ -23,6 +23,38 @@ from utils.file_utils import create_image_file class TestChatAssistantUpdate: + @pytest.mark.p2 + def test_update_rejects_non_dict_and_empty_llm_prompt(self, add_chat_assistants_func): + _, _, chat_assistants = add_chat_assistants_func + chat_assistant = chat_assistants[0] + + with pytest.raises(Exception) as exception_info: + chat_assistant.update.__wrapped__(chat_assistant, "bad") + assert "`update_message` must be a dict" in str(exception_info.value) + + with pytest.raises(Exception) as exception_info: + chat_assistant.update({"llm": {}}) + assert "`llm` cannot be empty" in str(exception_info.value) + + with pytest.raises(Exception) as exception_info: + chat_assistant.update({"prompt": {}}) + assert "`prompt` cannot be empty" in str(exception_info.value) + + @pytest.mark.p2 + def test_update_raises_on_nonzero_response(self, add_chat_assistants_func, monkeypatch): + _, _, chat_assistants = add_chat_assistants_func + chat_assistant = chat_assistants[0] + + class _DummyResponse: + def json(self): + return {"code": 1, "message": "boom"} + + monkeypatch.setattr(chat_assistant, "put", lambda *_args, **_kwargs: _DummyResponse()) + + with pytest.raises(Exception) as exception_info: + chat_assistant.update({"name": "error-case"}) + assert "boom" in str(exception_info.value) + @pytest.mark.parametrize( "payload, expected_message", [ diff --git a/test/testcases/test_sdk_api/test_chunk_management_within_dataset/test_list_chunks.py b/test/testcases/test_sdk_api/test_chunk_management_within_dataset/test_list_chunks.py index 4174d3fb14..6ecfbeb845 100644 --- a/test/testcases/test_sdk_api/test_chunk_management_within_dataset/test_list_chunks.py +++ b/test/testcases/test_sdk_api/test_chunk_management_within_dataset/test_list_chunks.py @@ -147,3 +147,14 @@ class TestChunksList: chunks = document.list_chunks() assert len(chunks) == 30, str(chunks) + + @pytest.mark.p2 + def test_list_chunks_invalid_document_id_raises(self, add_chunks): + _, document, _ = add_chunks + invalid_document = document.__class__( + document.rag, + {"id": "missing-document-id-for-chunks", "dataset_id": document.dataset_id}, + ) + with pytest.raises(Exception) as exception_info: + invalid_document.list_chunks() + assert str(exception_info.value), exception_info diff --git a/test/testcases/test_sdk_api/test_dataset_mangement/test_list_datasets.py b/test/testcases/test_sdk_api/test_dataset_mangement/test_list_datasets.py index c28366ba93..aa7e1b163b 100644 --- a/test/testcases/test_sdk_api/test_dataset_mangement/test_list_datasets.py +++ b/test/testcases/test_sdk_api/test_dataset_mangement/test_list_datasets.py @@ -218,6 +218,13 @@ class TestDatasetsList: client.list_datasets(**params) assert "lacks permission for dataset" in str(exception_info.value), str(exception_info.value) + @pytest.mark.p2 + def test_get_dataset_not_found_raises(self, client, monkeypatch): + monkeypatch.setattr(client, "list_datasets", lambda **_: []) + with pytest.raises(Exception) as exception_info: + client.get_dataset(name="missing-name-for-coverage") + assert "Dataset missing-name-for-coverage not found" in str(exception_info.value), str(exception_info.value) + @pytest.mark.p2 def test_name_empty(self, client): params = {"name": ""} diff --git a/test/testcases/test_sdk_api/test_file_management_within_dataset/test_download_document.py b/test/testcases/test_sdk_api/test_file_management_within_dataset/test_download_document.py index 3c9169fbb7..c6fad07ce1 100644 --- a/test/testcases/test_sdk_api/test_file_management_within_dataset/test_download_document.py +++ b/test/testcases/test_sdk_api/test_file_management_within_dataset/test_download_document.py @@ -64,6 +64,18 @@ class TestDocumentDownload: f.write(documents[0].download()) assert compare_by_hash(ragflow_tmp_dir / "ragflow_test_upload_0.txt", download_path), f"Downloaded file {i} does not match original" + @pytest.mark.p2 + def test_download_error_json_raises(self, add_documents): + dataset, documents = add_documents + document = documents[0] + invalid_document = document.__class__( + document.rag, + {"id": "missing-document-id-for-download", "dataset_id": dataset.id}, + ) + with pytest.raises(Exception) as exception_info: + invalid_document.download() + assert str(exception_info.value), exception_info + @pytest.mark.p3 def test_concurrent_download(add_dataset, tmp_path): diff --git a/test/testcases/test_sdk_api/test_file_management_within_dataset/test_parse_documents.py b/test/testcases/test_sdk_api/test_file_management_within_dataset/test_parse_documents.py index 3ff21178d4..97a9106628 100644 --- a/test/testcases/test_sdk_api/test_file_management_within_dataset/test_parse_documents.py +++ b/test/testcases/test_sdk_api/test_file_management_within_dataset/test_parse_documents.py @@ -14,10 +14,10 @@ # limitations under the License. # from concurrent.futures import ThreadPoolExecutor, as_completed - import pytest from common import bulk_upload_documents from ragflow_sdk import DataSet +from ragflow_sdk.modules.document import Document from utils import wait_for @@ -114,6 +114,116 @@ class TestDocumentsParse: validate_document_details(dataset, document_ids) +@pytest.mark.p2 +def test_get_documents_status_handles_retry_terminal_and_progress_paths(add_dataset_func, monkeypatch): + dataset = add_dataset_func + call_counts = {"doc-retry": 0, "doc-progress": 0, "doc-exception": 0} + + def _doc(doc_id, run, chunk_count, token_count, progress): + return Document( + dataset.rag, + { + "id": doc_id, + "dataset_id": dataset.id, + "run": run, + "chunk_count": chunk_count, + "token_count": token_count, + "progress": progress, + }, + ) + + def _list_documents(id=None, **_kwargs): + if id == "doc-retry": + call_counts["doc-retry"] += 1 + if call_counts["doc-retry"] == 1: + return [] + return [_doc("doc-retry", "DONE", 3, 5, 0.0)] + if id == "doc-progress": + call_counts["doc-progress"] += 1 + return [_doc("doc-progress", "RUNNING", 2, 4, 1.0)] + if id == "doc-exception": + call_counts["doc-exception"] += 1 + if call_counts["doc-exception"] == 1: + raise Exception("temporary list failure") + return [_doc("doc-exception", "DONE", 7, 11, 0.0)] + return [] + + monkeypatch.setattr(dataset, "list_documents", _list_documents) + monkeypatch.setattr("time.sleep", lambda *_args, **_kwargs: None) + + finished = dataset._get_documents_status(["doc-retry", "doc-progress", "doc-exception"]) + assert {item[0] for item in finished} == {"doc-retry", "doc-progress", "doc-exception"} + finished_map = {item[0]: item for item in finished} + assert finished_map["doc-retry"][1] == "DONE" + assert finished_map["doc-progress"][1] == "DONE" + assert finished_map["doc-exception"][1] == "DONE" + + +@pytest.mark.p2 +def test_parse_documents_keyboard_interrupt_triggers_cancel_then_returns_status(add_dataset_func, monkeypatch): + dataset = add_dataset_func + state = {"cancel_calls": 0, "status_calls": 0} + expected_status = [("doc-1", "DONE", 1, 2)] + + def _raise_keyboard_interrupt(_document_ids): + raise KeyboardInterrupt + + def _cancel(document_ids): + state["cancel_calls"] += 1 + assert document_ids == ["doc-1"] + + def _status(document_ids): + state["status_calls"] += 1 + assert document_ids == ["doc-1"] + return expected_status + + monkeypatch.setattr(dataset, "async_parse_documents", _raise_keyboard_interrupt) + monkeypatch.setattr(dataset, "async_cancel_parse_documents", _cancel) + monkeypatch.setattr(dataset, "_get_documents_status", _status) + + status = dataset.parse_documents(["doc-1"]) + assert status == expected_status + assert state["cancel_calls"] == 1 + assert state["status_calls"] == 1 + + +@pytest.mark.p2 +def test_parse_documents_happy_path_runs_initial_wait_then_returns_status(add_dataset_func, monkeypatch): + dataset = add_dataset_func + state = {"status_calls": 0} + + def _noop_parse(_document_ids): + return None + + def _status(document_ids): + state["status_calls"] += 1 + assert document_ids == ["doc-1"] + return [("doc-1", f"DONE-{state['status_calls']}", 1, 2)] + + monkeypatch.setattr(dataset, "async_parse_documents", _noop_parse) + monkeypatch.setattr(dataset, "_get_documents_status", _status) + + status = dataset.parse_documents(["doc-1"]) + assert state["status_calls"] == 2 + assert status == [("doc-1", "DONE-2", 1, 2)] + + +@pytest.mark.p2 +def test_async_cancel_parse_documents_raises_on_nonzero_code(add_dataset_func, monkeypatch): + dataset = add_dataset_func + + class _Resp: + @staticmethod + def json(): + return {"code": 102, "message": "cancel failed"} + + monkeypatch.setattr(dataset, "rm", lambda *_args, **_kwargs: _Resp()) + + with pytest.raises(Exception) as exc_info: + dataset.async_cancel_parse_documents(["doc-1"]) + assert "cancel failed" in str(exc_info.value), str(exc_info.value) + + @pytest.mark.p3 def test_parse_100_files(add_dataset_func, tmp_path): @wait_for(200, 1, "Document parsing timeout") diff --git a/test/testcases/test_sdk_api/test_file_management_within_dataset/test_update_document.py b/test/testcases/test_sdk_api/test_file_management_within_dataset/test_update_document.py index 00466ef338..69b84184a6 100644 --- a/test/testcases/test_sdk_api/test_file_management_within_dataset/test_update_document.py +++ b/test/testcases/test_sdk_api/test_file_management_within_dataset/test_update_document.py @@ -66,6 +66,14 @@ class TestDocumentsUpdated: else: document.update({"meta_fields": meta_fields}) + @pytest.mark.p2 + def test_meta_fields_invalid_type_guard_p2(self, add_documents): + _, documents = add_documents + document = documents[0] + with pytest.raises(Exception) as exception_info: + document.update({"meta_fields": "not-a-dict"}) + assert "meta_fields must be a dictionary" in str(exception_info.value), str(exception_info.value) + @pytest.mark.p2 @pytest.mark.parametrize( "chunk_method, expected_message", diff --git a/test/testcases/test_sdk_api/test_memory_management/test_list_memory.py b/test/testcases/test_sdk_api/test_memory_management/test_list_memory.py index 04cca63e7a..774cb59ccc 100644 --- a/test/testcases/test_sdk_api/test_memory_management/test_list_memory.py +++ b/test/testcases/test_sdk_api/test_memory_management/test_list_memory.py @@ -114,3 +114,13 @@ class TestMemoryList: "embd_id", "llm_id", "permissions", "description", "memory_size", "forgetting_policy", "temperature", "system_prompt", "user_prompt"]: assert hasattr(memory, field), memory_config + + @pytest.mark.p2 + def test_get_config_invalid_memory_id_raises(self, client): + memory_list = client.list_memory() + assert len(memory_list["memory_list"]) > 0, str(memory_list) + memory = memory_list["memory_list"][0] + memory.id = "missing-memory-id-for-config" + with pytest.raises(Exception) as exception_info: + memory.get_config() + assert str(exception_info.value), exception_info diff --git a/test/testcases/test_sdk_api/test_session_management/test_list_sessions_with_chat_assistant.py b/test/testcases/test_sdk_api/test_session_management/test_list_sessions_with_chat_assistant.py index 6889fd6ec3..df9c0e1c08 100644 --- a/test/testcases/test_sdk_api/test_session_management/test_list_sessions_with_chat_assistant.py +++ b/test/testcases/test_sdk_api/test_session_management/test_list_sessions_with_chat_assistant.py @@ -18,6 +18,20 @@ from concurrent.futures import ThreadPoolExecutor, as_completed class TestSessionsWithChatAssistantList: + @pytest.mark.p2 + def test_list_sessions_raises_on_nonzero_response(self, add_sessions_with_chat_assistant, monkeypatch): + chat_assistant, _ = add_sessions_with_chat_assistant + + class _DummyResponse: + def json(self): + return {"code": 1, "message": "boom"} + + monkeypatch.setattr(chat_assistant, "get", lambda *_args, **_kwargs: _DummyResponse()) + + with pytest.raises(Exception) as exception_info: + chat_assistant.list_sessions() + assert "boom" in str(exception_info.value) + @pytest.mark.p1 @pytest.mark.parametrize( "params, expected_page_size, expected_message", diff --git a/test/testcases/test_web_api/common.py b/test/testcases/test_web_api/common.py index cbbd1d768f..05e47eb186 100644 --- a/test/testcases/test_web_api/common.py +++ b/test/testcases/test_web_api/common.py @@ -292,7 +292,7 @@ def knowledge_graph(auth, dataset_id, params=None, *, headers=HEADERS): def delete_knowledge_graph(auth, dataset_id, payload=None, *, headers=HEADERS, data=None): - res = requests.delete(url=f"{HOST_ADDRESS}{KB_APP_URL}/{dataset_id}/delete_knowledge_graph", headers=headers, auth=auth, json=payload, data=data) + res = requests.delete(url=f"{HOST_ADDRESS}{KB_APP_URL}/{dataset_id}/knowledge_graph", headers=headers, auth=auth, json=payload, data=data) return res.json() @@ -434,6 +434,11 @@ def update_chunk(auth, payload=None, *, headers=HEADERS): return res.json() +def switch_chunks(auth, payload=None, *, headers=HEADERS): + res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/switch", headers=headers, auth=auth, json=payload) + return res.json() + + def delete_chunks(auth, payload=None, *, headers=HEADERS): res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/rm", headers=headers, auth=auth, json=payload) return res.json() diff --git a/test/testcases/test_web_api/test_agent_app/test_agents_webhook_unit.py b/test/testcases/test_web_api/test_agent_app/test_agents_webhook_unit.py new file mode 100644 index 0000000000..b58bbea553 --- /dev/null +++ b/test/testcases/test_web_api/test_agent_app/test_agents_webhook_unit.py @@ -0,0 +1,1208 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import asyncio +import base64 +import hashlib +import hmac +import importlib.util +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _Args(dict): + def get(self, key, default=None, type=None): + value = super().get(key, default) + if value is None or type is None: + return value + try: + return type(value) + except (TypeError, ValueError): + return default + + +class _DummyRequest: + def __init__( + self, + *, + path="/api/v1/webhook/agent-1", + method="POST", + headers=None, + content_length=0, + remote_addr="127.0.0.1", + args=None, + json_body=None, + raw_body=b"", + form=None, + files=None, + authorization=None, + ): + self.path = path + self.method = method + self.headers = headers or {} + self.content_length = content_length + self.remote_addr = remote_addr + self.args = args or {} + self.authorization = authorization + self.form = _AwaitableValue(form or {}) + self.files = _AwaitableValue(files or {}) + self._json_body = json_body + self._raw_body = raw_body + + async def get_json(self): + return self._json_body + + async def get_data(self): + return self._raw_body + + +class _CanvasRecord: + def __init__(self, *, canvas_category, dsl, user_id="tenant-1"): + self.canvas_category = canvas_category + self.dsl = dsl + self.user_id = user_id + + def to_dict(self): + return {"user_id": self.user_id, "dsl": self.dsl} + + +class _StubCanvas: + def __init__(self, dsl, user_id, agent_id, canvas_id=None): + self.dsl = dsl + self.user_id = user_id + self.agent_id = agent_id + self.canvas_id = canvas_id + + async def run(self, **_kwargs): + if False: + yield {} + + async def get_files_async(self, desc): + return {"files": desc} + + def __str__(self): + return "{}" + + +class _StubRedisConn: + def __init__(self): + self.bucket_result = [1] + self.bucket_exc = None + self.REDIS = object() + + def lua_token_bucket(self, **_kwargs): + if self.bucket_exc is not None: + raise self.bucket_exc + return self.bucket_result + + def get(self, _key): + return None + + def set_obj(self, _key, _obj, _ttl): + return None + + +def _run(coro): + return asyncio.run(coro) + + +def _default_webhook_params( + *, + security=None, + methods=None, + content_types="application/json", + schema=None, + execution_mode="Immediately", + response=None, +): + return { + "mode": "Webhook", + "methods": methods if methods is not None else ["POST"], + "security": security if security is not None else {}, + "content_types": content_types, + "schema": schema + if schema is not None + else { + "query": {"properties": {}, "required": []}, + "headers": {"properties": {}, "required": []}, + "body": {"properties": {}, "required": []}, + }, + "execution_mode": execution_mode, + "response": response if response is not None else {}, + } + + +def _make_webhook_cvs(module, *, params=None, dsl=None, canvas_category=None): + if dsl is None: + if params is None: + params = _default_webhook_params() + dsl = { + "components": { + "begin": { + "obj": {"component_name": "Begin", "params": params}, + "downstream": [], + "upstream": [], + } + } + } + if canvas_category is None: + canvas_category = module.CanvasCategory.Agent + return _CanvasRecord(canvas_category=canvas_category, dsl=dsl) + + +def _patch_background_task(monkeypatch, module): + def _fake_create_task(coro): + coro.close() + return None + + monkeypatch.setattr(module.asyncio, "create_task", _fake_create_task) + + +def _load_agents_app(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + agent_pkg = ModuleType("agent") + agent_pkg.__path__ = [] + canvas_mod = ModuleType("agent.canvas") + canvas_mod.Canvas = _StubCanvas + agent_pkg.canvas = canvas_mod + monkeypatch.setitem(sys.modules, "agent", agent_pkg) + monkeypatch.setitem(sys.modules, "agent.canvas", canvas_mod) + + services_pkg = ModuleType("api.db.services") + services_pkg.__path__ = [] + monkeypatch.setitem(sys.modules, "api.db.services", services_pkg) + + canvas_service_mod = ModuleType("api.db.services.canvas_service") + + class _StubUserCanvasService: + @staticmethod + def query(**_kwargs): + return [] + + @staticmethod + def get_list(*_args, **_kwargs): + return [] + + @staticmethod + def save(**_kwargs): + return True + + @staticmethod + def update_by_id(*_args, **_kwargs): + return True + + @staticmethod + def delete_by_id(*_args, **_kwargs): + return True + + @staticmethod + def get_by_id(_id): + return False, None + + canvas_service_mod.UserCanvasService = _StubUserCanvasService + monkeypatch.setitem(sys.modules, "api.db.services.canvas_service", canvas_service_mod) + services_pkg.canvas_service = canvas_service_mod + + file_service_mod = ModuleType("api.db.services.file_service") + + class _StubFileService: + @staticmethod + def upload_info(*_args, **_kwargs): + return {"id": "uploaded"} + + file_service_mod.FileService = _StubFileService + monkeypatch.setitem(sys.modules, "api.db.services.file_service", file_service_mod) + services_pkg.file_service = file_service_mod + + canvas_version_mod = ModuleType("api.db.services.user_canvas_version") + + class _StubUserCanvasVersionService: + @staticmethod + def insert(**_kwargs): + return True + + @staticmethod + def delete_all_versions(*_args, **_kwargs): + return True + + canvas_version_mod.UserCanvasVersionService = _StubUserCanvasVersionService + monkeypatch.setitem(sys.modules, "api.db.services.user_canvas_version", canvas_version_mod) + services_pkg.user_canvas_version = canvas_version_mod + + tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service") + + class _StubLLMFactoriesService: + @staticmethod + def get_api_key(*_args, **_kwargs): + return None + + tenant_llm_service_mod.LLMFactoriesService = _StubLLMFactoriesService + monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod) + services_pkg.tenant_llm_service = tenant_llm_service_mod + + redis_obj = _StubRedisConn() + redis_mod = ModuleType("rag.utils.redis_conn") + redis_mod.REDIS_CONN = redis_obj + monkeypatch.setitem(sys.modules, "rag.utils.redis_conn", redis_mod) + + module_path = repo_root / "api" / "apps" / "sdk" / "agents.py" + spec = importlib.util.spec_from_file_location("test_agents_webhook_unit", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + spec.loader.exec_module(module) + return module + + +def _assert_bad_request(res, expected_substring): + assert isinstance(res, tuple), res + payload, code = res + assert code == 400, res + assert payload["code"] == 400, payload + assert expected_substring in payload["message"], payload + + +@pytest.mark.p2 +def test_agents_crud_unit_branches(monkeypatch): + module = _load_agents_app(monkeypatch) + + monkeypatch.setattr( + module, + "request", + SimpleNamespace(args={"id": "missing", "title": "missing", "desc": "false", "page": "1", "page_size": "10"}), + ) + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: []) + res = module.list_agents.__wrapped__("tenant-1") + assert res["code"] == module.RetCode.DATA_ERROR + assert "doesn't exist" in res["message"] + + captured = {} + + def fake_get_list(_tenant_id, _page, _page_size, _orderby, desc, *_rest): + captured["desc"] = desc + return [{"id": "agent-1"}] + + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: [{"id": "agent-1"}]) + monkeypatch.setattr(module.UserCanvasService, "get_list", fake_get_list) + monkeypatch.setattr(module, "request", SimpleNamespace(args={"desc": "true"})) + res = module.list_agents.__wrapped__("tenant-1") + assert res["code"] == module.RetCode.SUCCESS + assert captured["desc"] is True + + async def req_no_dsl(): + return {"title": "agent-a"} + + monkeypatch.setattr(module, "get_request_json", req_no_dsl) + res = _run(module.create_agent.__wrapped__("tenant-1")) + assert res["code"] == module.RetCode.ARGUMENT_ERROR + assert "No DSL data in request" in res["message"] + + async def req_no_title(): + return {"dsl": {"components": {}}} + + monkeypatch.setattr(module, "get_request_json", req_no_title) + res = _run(module.create_agent.__wrapped__("tenant-1")) + assert res["code"] == module.RetCode.ARGUMENT_ERROR + assert "No title in request" in res["message"] + + async def req_dup(): + return {"dsl": {"components": {}}, "title": "agent-dup"} + + monkeypatch.setattr(module, "get_request_json", req_dup) + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: [object()]) + res = _run(module.create_agent.__wrapped__("tenant-1")) + assert res["code"] == module.RetCode.DATA_ERROR + assert "already exists" in res["message"] + + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module, "get_uuid", lambda: "agent-created") + monkeypatch.setattr(module.UserCanvasService, "save", lambda **_kwargs: False) + res = _run(module.create_agent.__wrapped__("tenant-1")) + assert res["code"] == module.RetCode.DATA_ERROR + assert "Fail to create agent" in res["message"] + + async def req_update(): + return {"dsl": {"nodes": []}, "title": " webhook-agent ", "unused": None} + + monkeypatch.setattr(module, "get_request_json", req_update) + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: False) + res = _run(module.update_agent.__wrapped__("tenant-1", "agent-1")) + assert res["code"] == module.RetCode.OPERATING_ERROR + + calls = {"update": 0, "insert": 0, "delete_versions": 0} + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: True) + monkeypatch.setattr( + module.UserCanvasService, + "update_by_id", + lambda *_args, **_kwargs: calls.__setitem__("update", calls["update"] + 1), + ) + monkeypatch.setattr( + module.UserCanvasVersionService, + "insert", + lambda **_kwargs: calls.__setitem__("insert", calls["insert"] + 1), + ) + monkeypatch.setattr( + module.UserCanvasVersionService, + "delete_all_versions", + lambda *_args, **_kwargs: calls.__setitem__("delete_versions", calls["delete_versions"] + 1), + ) + res = _run(module.update_agent.__wrapped__("tenant-1", "agent-1")) + assert res["code"] == module.RetCode.SUCCESS + assert calls == {"update": 1, "insert": 1, "delete_versions": 1} + + monkeypatch.setattr(module.UserCanvasService, "query", lambda **_kwargs: False) + res = module.delete_agent.__wrapped__("tenant-1", "agent-1") + assert res["code"] == module.RetCode.OPERATING_ERROR + + +@pytest.mark.p2 +def test_webhook_prechecks(monkeypatch): + module = _load_agents_app(monkeypatch) + monkeypatch.setattr(module, "request", _DummyRequest(headers={"Content-Type": "application/json"}, json_body={})) + + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (False, None)) + _assert_bad_request(_run(module.webhook("agent-1")), "Canvas not found") + + cvs = _make_webhook_cvs(module, canvas_category=module.CanvasCategory.DataFlow) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Dataflow can not be triggered") + + cvs = _make_webhook_cvs(module, dsl="invalid-dsl") + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Invalid DSL format") + + cvs = _make_webhook_cvs( + module, + dsl={"components": {"begin": {"obj": {"component_name": "Begin", "params": {"mode": "Chat"}}}}}, + ) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Webhook not configured") + + params = _default_webhook_params(methods=["GET"]) + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "not allowed") + + +@pytest.mark.p2 +def test_webhook_security_dispatch(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "application/json"}, json_body={}, args={"a": "b"}), + ) + + for security in ({}, {"auth_type": "none"}): + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=security)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id, _cvs=cvs: (True, _cvs)) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code"), res + assert res.status_code == 200 + + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security={"auth_type": "unsupported"})) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Unsupported auth_type") + + +@pytest.mark.p2 +def test_webhook_max_body_size(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + base_request = _DummyRequest(headers={"Content-Type": "application/json"}, json_body={}) + monkeypatch.setattr(module, "request", base_request) + + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security={"auth_type": "none"})) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code") + assert res.status_code == 200 + + security = {"auth_type": "none", "max_body_size": "123"} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=security)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Invalid max_body_size format") + + security = {"auth_type": "none", "max_body_size": "11mb"} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=security)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "exceeds maximum allowed size") + + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "application/json"}, json_body={}, content_length=2048), + ) + security = {"auth_type": "none", "max_body_size": "1kb"} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=security)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Request body too large") + + +@pytest.mark.p2 +def test_webhook_ip_whitelist(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "application/json"}, json_body={}, remote_addr="127.0.0.1"), + ) + + for whitelist in ([], ["127.0.0.0/24"], ["127.0.0.1"]): + security = {"auth_type": "none", "ip_whitelist": whitelist} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=security)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id, _cvs=cvs: (True, _cvs)) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code"), res + assert res.status_code == 200 + + security = {"auth_type": "none", "ip_whitelist": ["10.0.0.1"]} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=security)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "is not allowed") + + +@pytest.mark.p2 +def test_webhook_rate_limit(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + monkeypatch.setattr(module, "request", _DummyRequest(headers={"Content-Type": "application/json"}, json_body={})) + + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security={"auth_type": "none"})) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code") + assert res.status_code == 200 + + bad_limit = {"auth_type": "none", "rate_limit": {"limit": 0, "per": "minute"}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=bad_limit)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "rate_limit.limit must be > 0") + + bad_per = {"auth_type": "none", "rate_limit": {"limit": 1, "per": "week"}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=bad_per)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Invalid rate_limit.per") + + module.REDIS_CONN.bucket_result = [0] + module.REDIS_CONN.bucket_exc = None + denied = {"auth_type": "none", "rate_limit": {"limit": 1, "per": "minute"}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=denied)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Too many requests") + + module.REDIS_CONN.bucket_result = [1] + module.REDIS_CONN.bucket_exc = RuntimeError("redis failure") + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=denied)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Rate limit error") + + +@pytest.mark.p2 +def test_webhook_token_basic_jwt_auth(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + monkeypatch.setattr(module, "request", _DummyRequest(headers={"Content-Type": "application/json"}, json_body={})) + + token_security = {"auth_type": "token", "token": {"token_header": "X-TOKEN", "token_value": "ok"}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=token_security)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Invalid token authentication") + + monkeypatch.setattr( + module, + "request", + _DummyRequest( + headers={"Content-Type": "application/json"}, + json_body={}, + authorization=SimpleNamespace(username="u", password="bad"), + ), + ) + basic_security = {"auth_type": "basic", "basic_auth": {"username": "u", "password": "p"}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=basic_security)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Invalid Basic Auth credentials") + + monkeypatch.setattr(module, "request", _DummyRequest(headers={"Content-Type": "application/json"}, json_body={})) + jwt_missing_secret = {"auth_type": "jwt", "jwt": {}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=jwt_missing_secret)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "JWT secret not configured") + + jwt_base = {"auth_type": "jwt", "jwt": {"secret": "secret"}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=jwt_base)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Missing Bearer token") + + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "application/json", "Authorization": "Bearer "}, json_body={}), + ) + _assert_bad_request(_run(module.webhook("agent-1")), "Empty Bearer token") + + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "application/json", "Authorization": "Bearer token"}, json_body={}), + ) + monkeypatch.setattr(module.jwt, "decode", lambda *_args, **_kwargs: (_ for _ in ()).throw(Exception("decode boom"))) + _assert_bad_request(_run(module.webhook("agent-1")), "Invalid JWT") + + monkeypatch.setattr(module.jwt, "decode", lambda *_args, **_kwargs: {"exp": 1}) + jwt_reserved = {"auth_type": "jwt", "jwt": {"secret": "secret", "required_claims": ["exp"]}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=jwt_reserved)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Reserved JWT claim cannot be required") + + monkeypatch.setattr(module.jwt, "decode", lambda *_args, **_kwargs: {}) + jwt_missing_claim = {"auth_type": "jwt", "jwt": {"secret": "secret", "required_claims": ["role"]}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=jwt_missing_claim)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + _assert_bad_request(_run(module.webhook("agent-1")), "Missing JWT claim") + + captured = {} + + def fake_decode(token, options, **kwargs): + captured["token"] = token + captured["options"] = options + captured["kwargs"] = kwargs + return {"role": "admin"} + + monkeypatch.setattr(module.jwt, "decode", fake_decode) + jwt_success = { + "auth_type": "jwt", + "jwt": { + "secret": "secret", + "audience": "aud", + "issuer": "iss", + "required_claims": "role", + }, + } + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=jwt_success)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code") + assert res.status_code == 200 + assert captured["kwargs"]["audience"] == "aud" + assert captured["kwargs"]["issuer"] == "iss" + assert captured["options"]["verify_aud"] is True + assert captured["options"]["verify_iss"] is True + + monkeypatch.setattr(module.jwt, "decode", lambda *_args, **_kwargs: {}) + jwt_success_invalid_type = {"auth_type": "jwt", "jwt": {"secret": "secret", "required_claims": 123}} + cvs = _make_webhook_cvs(module, params=_default_webhook_params(security=jwt_success_invalid_type)) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code") + assert res.status_code == 200 + + +@pytest.mark.p2 +def test_webhook_parse_request_branches(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + security = {"auth_type": "none"} + params = _default_webhook_params(security=security, content_types="application/json") + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "text/plain"}, raw_body=b'{"x":1}', json_body={}), + ) + with pytest.raises(ValueError, match="Invalid Content-Type"): + _run(module.webhook("agent-1")) + + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "application/json"}, json_body={"x": 1}, args={"q": "1"}), + ) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code") + assert res.status_code == 200 + + params = _default_webhook_params(security=security, content_types="multipart/form-data") + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + files = {f"file{i}": object() for i in range(11)} + monkeypatch.setattr( + module, + "request", + _DummyRequest( + headers={"Content-Type": "multipart/form-data"}, + form={"key": "value"}, + files=files, + json_body={}, + ), + ) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code") + assert res.status_code == 200 + + uploaded = {"count": 0} + monkeypatch.setattr( + module.FileService, + "upload_info", + lambda *_args, **_kwargs: uploaded.__setitem__("count", uploaded["count"] + 1) or {"id": "uploaded"}, + ) + monkeypatch.setattr( + module, + "request", + _DummyRequest( + headers={"Content-Type": "multipart/form-data"}, + form={"k": "v"}, + files={"file1": object()}, + json_body={}, + ), + ) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code") + assert res.status_code == 200 + assert uploaded["count"] == 1 + + +@pytest.mark.p2 +def test_webhook_canvas_constructor_exception(monkeypatch): + module = _load_agents_app(monkeypatch) + + params = _default_webhook_params(security={"auth_type": "none"}) + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "application/json"}, json_body={}), + ) + monkeypatch.setattr(module, "Canvas", lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("canvas init failed"))) + + def fake_error_result(*, code, message): + return SimpleNamespace(code=code, message=message) + + monkeypatch.setattr(module, "get_data_error_result", fake_error_result) + res = _run(module.webhook("agent-1")) + assert isinstance(res, SimpleNamespace) + assert res.code == module.RetCode.BAD_REQUEST + assert "canvas init failed" in res.message + assert res.status_code == module.RetCode.BAD_REQUEST + + +@pytest.mark.p2 +def test_webhook_trace_polling_branches(monkeypatch): + module = _load_agents_app(monkeypatch) + + # Missing since_ts. + monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args())) + res = _run(module.webhook_trace("agent-1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["webhook_id"] is None + assert res["data"]["events"] == [] + assert res["data"]["finished"] is False + + # since_ts provided but no Redis data. + monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args({"since_ts": "100.0"}))) + monkeypatch.setattr(module.REDIS_CONN, "get", lambda _k: None) + res = _run(module.webhook_trace("agent-1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["webhook_id"] is None + assert res["data"]["next_since_ts"] == 100.0 + assert res["data"]["events"] == [] + assert res["data"]["finished"] is False + + webhooks_obj = { + "webhooks": { + "101.0": { + "events": [ + {"event": "message", "ts": 101.2, "data": {"content": "a"}}, + {"event": "finished", "ts": 102.5}, + ] + }, + "99.0": {"events": [{"event": "message", "ts": 99.1}]}, + } + } + raw = json.dumps(webhooks_obj) + monkeypatch.setattr(module.REDIS_CONN, "get", lambda _k: raw) + + # No candidates newer than since_ts. + monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args({"since_ts": "200.0"}))) + res = _run(module.webhook_trace("agent-1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["webhook_id"] is None + assert res["data"]["next_since_ts"] == 200.0 + assert res["data"]["events"] == [] + assert res["data"]["finished"] is False + + # Candidate exists and webhook id is assigned. + monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args({"since_ts": "100.0"}))) + res = _run(module.webhook_trace("agent-1")) + assert res["code"] == module.RetCode.SUCCESS + webhook_id = res["data"]["webhook_id"] + assert webhook_id + assert res["data"]["events"] == [] + assert res["data"]["next_since_ts"] == 101.0 + assert res["data"]["finished"] is False + + # Invalid webhook id. + monkeypatch.setattr( + module, + "request", + SimpleNamespace(args=_Args({"since_ts": "100.0", "webhook_id": "bad-id"})), + ) + res = _run(module.webhook_trace("agent-1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["webhook_id"] == "bad-id" + assert res["data"]["events"] == [] + assert res["data"]["next_since_ts"] == 100.0 + assert res["data"]["finished"] is True + + # Valid webhook id with event filtering and finished flag. + monkeypatch.setattr( + module, + "request", + SimpleNamespace(args=_Args({"since_ts": "101.0", "webhook_id": webhook_id})), + ) + res = _run(module.webhook_trace("agent-1")) + assert res["code"] == module.RetCode.SUCCESS + assert res["data"]["webhook_id"] == webhook_id + assert [event["ts"] for event in res["data"]["events"]] == [101.2, 102.5] + assert res["data"]["next_since_ts"] == 102.5 + assert res["data"]["finished"] is True + + +@pytest.mark.p2 +def test_webhook_parse_request_form_and_raw_body_paths(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + security = {"auth_type": "none"} + + def _run_with(params, req): + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id, _cvs=cvs: (True, _cvs)) + monkeypatch.setattr(module, "request", req) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code"), res + assert res.status_code == 200 + + _run_with( + _default_webhook_params(security=security, content_types="application/x-www-form-urlencoded"), + _DummyRequest( + headers={"Content-Type": "application/x-www-form-urlencoded"}, + form={"a": "1", "b": "2"}, + json_body={}, + ), + ) + + _run_with( + _default_webhook_params(security=security, content_types="text/plain"), + _DummyRequest(headers={"Content-Type": "text/plain"}, raw_body=b'{"k": 1}', json_body={}), + ) + + _run_with( + _default_webhook_params(security=security, content_types="text/plain"), + _DummyRequest(headers={"Content-Type": "text/plain"}, raw_body=b"{bad-json}", json_body={}), + ) + + _run_with( + _default_webhook_params(security=security, content_types="text/plain"), + _DummyRequest(headers={"Content-Type": "text/plain"}, raw_body=b"", json_body={}), + ) + + class _BrokenRawRequest(_DummyRequest): + async def get_data(self): + raise RuntimeError("raw read failed") + + _run_with( + _default_webhook_params(security=security, content_types="text/plain"), + _BrokenRawRequest(headers={"Content-Type": "text/plain"}, json_body={}), + ) + + +@pytest.mark.p2 +def test_webhook_schema_extract_cast_defaults_and_validation_errors(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + base_schema = { + "query": { + "properties": { + "q_file": {"type": "file"}, + "q_object": {"type": "object"}, + "q_boolean": {"type": "boolean"}, + "q_number": {"type": "number"}, + "q_string": {"type": "string"}, + "q_array": {"type": "array"}, + "q_null": {"type": "null"}, + "q_default_none": {}, + }, + "required": [], + }, + "headers": {"properties": {"Content-Type": {"type": "string"}}, "required": []}, + "body": { + "properties": { + "bool_true": {"type": "boolean"}, + "bool_false": {"type": "boolean"}, + "number_int": {"type": "number"}, + "number_float": {"type": "number"}, + "obj": {"type": "object"}, + "arr": {"type": "array"}, + "text": {"type": "string"}, + "file_list": {"type": "file"}, + "unknown": {"type": "mystery"}, + }, + "required": [ + "bool_true", + "number_int", + "obj", + "arr", + "text", + "file_list", + "unknown", + ], + }, + } + + params = _default_webhook_params( + security={"auth_type": "none"}, + content_types="application/json", + schema=base_schema, + ) + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + monkeypatch.setattr( + module, + "request", + _DummyRequest( + headers={"Content-Type": "application/json"}, + args={}, + json_body={ + "bool_true": "true", + "bool_false": "0", + "number_int": "-3", + "number_float": "2.5", + "obj": '{"a": 1}', + "arr": "[1, 2]", + "text": "hello", + "file_list": ["f1"], + "unknown": "mystery", + }, + ), + ) + res = _run(module.webhook("agent-1")) + assert hasattr(res, "status_code"), res + assert res.status_code == 200 + + failure_cases = [ + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"must": {"type": "string"}}, "required": ["must"]}}, + {}, + "missing required field", + ), + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"flag": {"type": "boolean"}}, "required": ["flag"]}}, + {"flag": "maybe"}, + "auto-cast failed", + ), + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"num": {"type": "number"}}, "required": ["num"]}}, + {"num": "abc"}, + "auto-cast failed", + ), + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"obj": {"type": "object"}}, "required": ["obj"]}}, + {"obj": "[]"}, + "auto-cast failed", + ), + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"arr": {"type": "array"}}, "required": ["arr"]}}, + {"arr": "{}"}, + "auto-cast failed", + ), + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"num": {"type": "number"}}, "required": ["num"]}}, + {"num": []}, + "type mismatch", + ), + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"arr": {"type": "array"}}, "required": ["arr"]}}, + {"arr": 3}, + "type mismatch", + ), + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"arr": {"type": "array"}}, "required": ["arr"]}}, + {"arr": [1, "x"]}, + "type mismatch", + ), + ( + {"query": {"properties": {}, "required": []}, "headers": {"properties": {}, "required": []}, "body": {"properties": {"file": {"type": "file"}}, "required": ["file"]}}, + {"file": "inline-file"}, + "type mismatch", + ), + ] + + for schema, body_payload, expected_substring in failure_cases: + params = _default_webhook_params( + security={"auth_type": "none"}, + content_types="application/json", + schema=schema, + ) + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id, _cvs=cvs: (True, _cvs)) + monkeypatch.setattr( + module, + "request", + _DummyRequest(headers={"Content-Type": "application/json"}, json_body=body_payload), + ) + res = _run(module.webhook("agent-1")) + _assert_bad_request(res, expected_substring) + + +@pytest.mark.p2 +def test_webhook_immediate_response_status_and_template_validation(monkeypatch): + module = _load_agents_app(monkeypatch) + _patch_background_task(monkeypatch, module) + + def _run_case(response_cfg): + params = _default_webhook_params( + security={"auth_type": "none"}, + content_types="application/json", + response=response_cfg, + ) + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id, _cvs=cvs: (True, _cvs)) + monkeypatch.setattr(module, "request", _DummyRequest(headers={"Content-Type": "application/json"}, json_body={})) + return _run(module.webhook("agent-1")) + + _assert_bad_request(_run_case({"status": "abc"}), "Invalid response status code") + _assert_bad_request(_run_case({"status": 500}), "must be between 200 and 399") + + empty_res = _run_case({"status": 204, "body_template": ""}) + assert empty_res.status_code == 204 + assert empty_res.content_type == "application/json" + assert _run(empty_res.get_data(as_text=True)) == "null" + + json_res = _run_case({"status": 201, "body_template": '{"ok": true}'}) + assert json_res.status_code == 201 + assert json_res.content_type == "application/json" + assert json.loads(_run(json_res.get_data(as_text=True))) == {"ok": True} + + plain_res = _run_case({"status": 202, "body_template": "plain-text"}) + assert plain_res.status_code == 202 + assert plain_res.content_type == "text/plain" + assert _run(plain_res.get_data(as_text=True)) == "plain-text" + + +@pytest.mark.p2 +def test_webhook_background_run_success_and_error_trace_paths(monkeypatch): + module = _load_agents_app(monkeypatch) + + redis_store = {} + + def redis_get(key): + return redis_store.get(key) + + def redis_set_obj(key, obj, _ttl): + redis_store[key] = json.dumps(obj) + + monkeypatch.setattr(module.REDIS_CONN, "get", redis_get) + monkeypatch.setattr(module.REDIS_CONN, "set_obj", redis_set_obj) + + update_calls = [] + monkeypatch.setattr(module.UserCanvasService, "update_by_id", lambda *_args, **_kwargs: update_calls.append(True)) + + tasks = [] + + def _capture_task(coro): + tasks.append(coro) + return SimpleNamespace() + + monkeypatch.setattr(module.asyncio, "create_task", _capture_task) + + class _CanvasSuccess(_StubCanvas): + async def run(self, **_kwargs): + yield {"event": "message", "data": {"content": "ok"}} + + def __str__(self): + return "{}" + + monkeypatch.setattr(module, "Canvas", _CanvasSuccess) + + params = _default_webhook_params(security={"auth_type": "none"}, content_types="application/json") + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + monkeypatch.setattr( + module, + "request", + _DummyRequest(path="/api/v1/webhook_test/agent-1", headers={"Content-Type": "application/json"}, json_body={}), + ) + + res = _run(module.webhook("agent-1")) + assert res.status_code == 200 + assert len(tasks) == 1 + _run(tasks.pop(0)) + assert update_calls == [True] + + key = "webhook-trace-agent-1-logs" + trace_obj = json.loads(redis_store[key]) + ws = next(iter(trace_obj["webhooks"].values())) + events = ws["events"] + assert any(event.get("event") == "message" for event in events) + assert any(event.get("event") == "finished" and event.get("success") is True for event in events) + + class _CanvasError(_StubCanvas): + async def run(self, **_kwargs): + raise RuntimeError("run failed") + yield {} + + monkeypatch.setattr(module, "Canvas", _CanvasError) + tasks.clear() + redis_store.clear() + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id, _cvs=cvs: (True, _cvs)) + res = _run(module.webhook("agent-1")) + assert res.status_code == 200 + _run(tasks.pop(0)) + trace_obj = json.loads(redis_store[key]) + ws = next(iter(trace_obj["webhooks"].values())) + events = ws["events"] + assert any(event.get("event") == "error" for event in events) + assert any(event.get("event") == "finished" and event.get("success") is False for event in events) + + log_messages = [] + monkeypatch.setattr(module.logging, "exception", lambda msg, *_args, **_kwargs: log_messages.append(str(msg))) + monkeypatch.setattr(module.REDIS_CONN, "get", lambda _key: "{") + monkeypatch.setattr(module.REDIS_CONN, "set_obj", lambda *_args, **_kwargs: None) + tasks.clear() + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id, _cvs=cvs: (True, _cvs)) + _run(module.webhook("agent-1")) + _run(tasks.pop(0)) + assert any("Failed to append webhook trace" in msg for msg in log_messages) + + +@pytest.mark.p2 +def test_webhook_sse_success_and_exception_paths(monkeypatch): + module = _load_agents_app(monkeypatch) + + redis_store = {} + monkeypatch.setattr(module.REDIS_CONN, "get", lambda key: redis_store.get(key)) + monkeypatch.setattr(module.REDIS_CONN, "set_obj", lambda key, obj, _ttl: redis_store.__setitem__(key, json.dumps(obj))) + + params = _default_webhook_params( + security={"auth_type": "none"}, + content_types="application/json", + execution_mode="Deferred", + ) + cvs = _make_webhook_cvs(module, params=params) + monkeypatch.setattr(module.UserCanvasService, "get_by_id", lambda _id: (True, cvs)) + + class _CanvasSSESuccess(_StubCanvas): + async def run(self, **_kwargs): + yield {"event": "message", "data": {"content": "x", "start_to_think": True}} + yield {"event": "message", "data": {"content": "y", "end_to_think": True}} + yield {"event": "message", "data": {"content": "Hello"}} + yield {"event": "message_end", "data": {"status": "201"}} + + monkeypatch.setattr(module, "Canvas", _CanvasSSESuccess) + monkeypatch.setattr( + module, + "request", + _DummyRequest(path="/api/v1/webhook_test/agent-1", headers={"Content-Type": "application/json"}, json_body={}), + ) + res = _run(module.webhook("agent-1")) + assert res.status_code == 201 + payload = json.loads(_run(res.get_data(as_text=True))) + assert payload == {"message": "Hello", "success": True, "code": 201} + + class _CanvasSSEError(_StubCanvas): + async def run(self, **_kwargs): + raise RuntimeError("sse failed") + yield {} + + monkeypatch.setattr(module, "Canvas", _CanvasSSEError) + monkeypatch.setattr( + module, + "request", + _DummyRequest(path="/api/v1/webhook_test/agent-1", headers={"Content-Type": "application/json"}, json_body={}), + ) + res = _run(module.webhook("agent-1")) + assert res.status_code == 400 + payload = json.loads(_run(res.get_data(as_text=True))) + assert payload["code"] == 400 + assert payload["success"] is False + assert "sse failed" in payload["message"] + + +@pytest.mark.p2 +def test_webhook_trace_encoded_id_generation(monkeypatch): + module = _load_agents_app(monkeypatch) + + webhooks_obj = { + "webhooks": { + "101.0": { + "events": [{"event": "message", "ts": 101.2}], + } + } + } + monkeypatch.setattr(module.REDIS_CONN, "get", lambda _key: json.dumps(webhooks_obj)) + monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args({"since_ts": "100.0"}))) + res = _run(module.webhook_trace("agent-1")) + assert res["code"] == module.RetCode.SUCCESS + + expected = base64.urlsafe_b64encode( + hmac.new( + b"webhook_id_secret", + b"101.0", + hashlib.sha256, + ).digest() + ).decode("utf-8").rstrip("=") + assert res["data"]["webhook_id"] == expected diff --git a/test/testcases/test_web_api/test_api_app/test_api_tokens_unit.py b/test/testcases/test_web_api/test_api_app/test_api_tokens_unit.py new file mode 100644 index 0000000000..b5c3d56546 --- /dev/null +++ b/test/testcases/test_web_api/test_api_app/test_api_tokens_unit.py @@ -0,0 +1,247 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _ExprField: + def __init__(self, name): + self.name = name + + def __eq__(self, other): + return (self.name, other) + + +class _DummyAPITokenModel: + tenant_id = _ExprField("tenant_id") + token = _ExprField("token") + + +def _run(coro): + return asyncio.run(coro) + + +def _load_api_app(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + quart_mod = ModuleType("quart") + quart_mod.request = SimpleNamespace(args={}) + monkeypatch.setitem(sys.modules, "quart", quart_mod) + + apps_mod = ModuleType("api.apps") + apps_mod.__path__ = [str(repo_root / "api" / "apps")] + apps_mod.login_required = lambda fn: fn + apps_mod.current_user = SimpleNamespace(id="user-1") + monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + + api_utils_mod = ModuleType("api.utils.api_utils") + + async def _get_request_json(): + return {} + + api_utils_mod.generate_confirmation_token = lambda: "token-123" + api_utils_mod.get_request_json = _get_request_json + api_utils_mod.get_json_result = lambda data=None, message="", code=0: { + "code": code, + "message": message, + "data": data, + } + api_utils_mod.get_data_error_result = lambda message="", code=400, data=None: { + "code": code, + "message": message, + "data": data, + } + api_utils_mod.server_error_response = lambda exc: { + "code": 500, + "message": str(exc), + "data": None, + } + api_utils_mod.validate_request = lambda *_args, **_kwargs: (lambda fn: fn) + monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod) + + api_service_mod = ModuleType("api.db.services.api_service") + + class _StubAPITokenService: + @staticmethod + def save(**_kwargs): + return True + + @staticmethod + def query(**_kwargs): + return [] + + @staticmethod + def filter_delete(_conds): + return True + + class _StubAPI4ConversationService: + @staticmethod + def stats(*_args, **_kwargs): + return [] + + api_service_mod.APITokenService = _StubAPITokenService + api_service_mod.API4ConversationService = _StubAPI4ConversationService + monkeypatch.setitem(sys.modules, "api.db.services.api_service", api_service_mod) + + user_service_mod = ModuleType("api.db.services.user_service") + + class _StubUserTenantService: + @staticmethod + def query(**_kwargs): + return [SimpleNamespace(tenant_id="tenant-1")] + + user_service_mod.UserTenantService = _StubUserTenantService + monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod) + + db_models_mod = ModuleType("api.db.db_models") + db_models_mod.APIToken = _DummyAPITokenModel + monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod) + + time_utils_mod = ModuleType("common.time_utils") + time_utils_mod.current_timestamp = lambda: 123 + time_utils_mod.datetime_format = lambda _dt: "2026-01-01 00:00:00" + monkeypatch.setitem(sys.modules, "common.time_utils", time_utils_mod) + + module_path = repo_root / "api" / "apps" / "api_app.py" + spec = importlib.util.spec_from_file_location("test_api_tokens_unit_module", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + spec.loader.exec_module(module) + return module + + +@pytest.mark.p2 +def test_new_token_branches_and_error_paths(monkeypatch): + module = _load_api_app(monkeypatch) + + async def req_canvas(): + return {"canvas_id": "canvas-1"} + + monkeypatch.setattr(module, "get_request_json", req_canvas) + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: []) + res = _run(module.new_token()) + assert res["message"] == "Tenant not found!" + + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module.APITokenService, "save", lambda **_kwargs: True) + res = _run(module.new_token()) + assert res["code"] == 0 + assert res["data"]["tenant_id"] == "tenant-1" + assert res["data"]["dialog_id"] == "canvas-1" + assert res["data"]["source"] == "agent" + + monkeypatch.setattr(module.APITokenService, "save", lambda **_kwargs: False) + res = _run(module.new_token()) + assert res["message"] == "Fail to new a dialog!" + + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("query failed"))) + res = _run(module.new_token()) + assert res["code"] == 500 + assert "query failed" in res["message"] + + +@pytest.mark.p2 +def test_token_list_tenant_guard_and_exception(monkeypatch): + module = _load_api_app(monkeypatch) + + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module, "request", SimpleNamespace(args={"dialog_id": "d1"})) + res = module.token_list() + assert res["message"] == "Tenant not found!" + + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module, "request", SimpleNamespace(args={})) + res = module.token_list() + assert res["code"] == 500 + assert "canvas_id" in res["message"] + + +@pytest.mark.p2 +def test_rm_exception_path(monkeypatch): + module = _load_api_app(monkeypatch) + + async def req_rm(): + return {"tokens": ["tok-1"], "tenant_id": "tenant-1"} + + monkeypatch.setattr(module, "get_request_json", req_rm) + monkeypatch.setattr( + module.APITokenService, + "filter_delete", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("delete failed")), + ) + + res = _run(module.rm()) + assert res["code"] == 500 + assert "delete failed" in res["message"] + + +@pytest.mark.p2 +def test_stats_aggregation_and_error_paths(monkeypatch): + module = _load_api_app(monkeypatch) + + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module, "request", SimpleNamespace(args={})) + res = module.stats() + assert res["message"] == "Tenant not found!" + + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module, "request", SimpleNamespace(args={"canvas_id": "canvas-1"})) + monkeypatch.setattr( + module.API4ConversationService, + "stats", + lambda *_args, **_kwargs: [ + { + "dt": "2026-01-01", + "pv": 3, + "uv": 2, + "tokens": 100, + "duration": 9.9, + "round": 1, + "thumb_up": 0, + } + ], + ) + res = module.stats() + assert res["code"] == 0 + assert res["data"]["pv"] == [("2026-01-01", 3)] + assert res["data"]["uv"] == [("2026-01-01", 2)] + assert res["data"]["round"] == [("2026-01-01", 1)] + assert res["data"]["thumb_up"] == [("2026-01-01", 0)] + assert res["data"]["tokens"] == [("2026-01-01", 0.1)] + assert res["data"]["speed"] == [("2026-01-01", 10.0)] + + monkeypatch.setattr( + module.API4ConversationService, + "stats", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("stats failed")), + ) + res = module.stats() + assert res["code"] == 500 + assert "stats failed" in res["message"] diff --git a/test/testcases/test_web_api/test_auth_app/test_oidc_client_unit.py b/test/testcases/test_web_api/test_auth_app/test_oidc_client_unit.py new file mode 100644 index 0000000000..f1e620d65d --- /dev/null +++ b/test/testcases/test_web_api/test_auth_app/test_oidc_client_unit.py @@ -0,0 +1,484 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _FakeResponse: + def __init__(self, payload=None, err=None): + self._payload = payload or {} + self._err = err + + def raise_for_status(self): + if self._err: + raise self._err + + def json(self): + return self._payload + + +class _DummyJwkClient: + def __init__(self, _jwks_uri): + self._key = "dummy-signing-key" + + def get_signing_key_from_jwt(self, _id_token): + return SimpleNamespace(key=self._key) + + +def _load_auth_modules(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + api_pkg = ModuleType("api") + api_pkg.__path__ = [str(repo_root / "api")] + apps_pkg = ModuleType("api.apps") + apps_pkg.__path__ = [str(repo_root / "api" / "apps")] + auth_pkg = ModuleType("api.apps.auth") + auth_pkg.__path__ = [str(repo_root / "api" / "apps" / "auth")] + + monkeypatch.setitem(sys.modules, "api", api_pkg) + monkeypatch.setitem(sys.modules, "api.apps", apps_pkg) + monkeypatch.setitem(sys.modules, "api.apps.auth", auth_pkg) + + for mod_name in ["api.apps.auth.oauth", "api.apps.auth.oidc"]: + sys.modules.pop(mod_name, None) + + oauth_path = repo_root / "api" / "apps" / "auth" / "oauth.py" + oauth_spec = importlib.util.spec_from_file_location("api.apps.auth.oauth", oauth_path) + oauth_module = importlib.util.module_from_spec(oauth_spec) + monkeypatch.setitem(sys.modules, "api.apps.auth.oauth", oauth_module) + oauth_spec.loader.exec_module(oauth_module) + + oidc_path = repo_root / "api" / "apps" / "auth" / "oidc.py" + oidc_spec = importlib.util.spec_from_file_location("api.apps.auth.oidc", oidc_path) + oidc_module = importlib.util.module_from_spec(oidc_spec) + monkeypatch.setitem(sys.modules, "api.apps.auth.oidc", oidc_module) + oidc_spec.loader.exec_module(oidc_module) + + return oauth_module, oidc_module + + +def _load_github_module(monkeypatch): + _load_auth_modules(monkeypatch) + repo_root = Path(__file__).resolve().parents[4] + + sys.modules.pop("api.apps.auth.github", None) + github_path = repo_root / "api" / "apps" / "auth" / "github.py" + github_spec = importlib.util.spec_from_file_location("api.apps.auth.github", github_path) + github_module = importlib.util.module_from_spec(github_spec) + monkeypatch.setitem(sys.modules, "api.apps.auth.github", github_module) + github_spec.loader.exec_module(github_module) + return github_module + + +def _load_auth_init_module(monkeypatch): + _load_auth_modules(monkeypatch) + repo_root = Path(__file__).resolve().parents[4] + + github_mod = ModuleType("api.apps.auth.github") + + class _StubGithubOAuthClient: + def __init__(self, config): + self.config = config + + github_mod.GithubOAuthClient = _StubGithubOAuthClient + monkeypatch.setitem(sys.modules, "api.apps.auth.github", github_mod) + + init_path = repo_root / "api" / "apps" / "auth" / "__init__.py" + init_spec = importlib.util.spec_from_file_location( + "api.apps.auth", + init_path, + submodule_search_locations=[str(repo_root / "api" / "apps" / "auth")], + ) + init_module = importlib.util.module_from_spec(init_spec) + monkeypatch.setitem(sys.modules, "api.apps.auth", init_module) + init_spec.loader.exec_module(init_module) + return init_module + + +def _base_config(): + return { + "issuer": "https://issuer.example", + "client_id": "client-1", + "client_secret": "secret-1", + "redirect_uri": "https://app.example/callback", + } + + +def _metadata(issuer): + return { + "issuer": issuer, + "jwks_uri": f"{issuer}/jwks", + "authorization_endpoint": f"{issuer}/authorize", + "token_endpoint": f"{issuer}/token", + "userinfo_endpoint": f"{issuer}/userinfo", + } + + +def _make_client(monkeypatch, oidc_module): + monkeypatch.setattr(oidc_module.OIDCClient, "_load_oidc_metadata", staticmethod(lambda issuer: _metadata(issuer))) + return oidc_module.OIDCClient(_base_config()) + + +@pytest.mark.p2 +def test_oidc_init_requires_issuer(monkeypatch): + _, oidc_module = _load_auth_modules(monkeypatch) + + with pytest.raises(ValueError) as exc_info: + oidc_module.OIDCClient({"client_id": "cid"}) + + assert str(exc_info.value) == "Missing issuer in configuration." + + +@pytest.mark.p2 +def test_oidc_init_loads_metadata_and_sets_endpoints(monkeypatch): + _, oidc_module = _load_auth_modules(monkeypatch) + monkeypatch.setattr(oidc_module.OIDCClient, "_load_oidc_metadata", staticmethod(lambda issuer: _metadata(issuer))) + + client = oidc_module.OIDCClient(_base_config()) + + assert client.issuer == "https://issuer.example" + assert client.jwks_uri == "https://issuer.example/jwks" + assert client.authorization_url == "https://issuer.example/authorize" + assert client.token_url == "https://issuer.example/token" + assert client.userinfo_url == "https://issuer.example/userinfo" + + +@pytest.mark.p2 +def test_load_oidc_metadata_success_and_wraps_failure(monkeypatch): + _, oidc_module = _load_auth_modules(monkeypatch) + + calls = {} + + def _ok_sync_request(method, url, timeout): + calls.update({"method": method, "url": url, "timeout": timeout}) + return _FakeResponse(_metadata("https://issuer.example")) + + monkeypatch.setattr(oidc_module, "sync_request", _ok_sync_request) + metadata = oidc_module.OIDCClient._load_oidc_metadata("https://issuer.example") + assert metadata["jwks_uri"] == "https://issuer.example/jwks" + assert calls == { + "method": "GET", + "url": "https://issuer.example/.well-known/openid-configuration", + "timeout": 7, + } + + def _boom_sync_request(*_args, **_kwargs): + raise RuntimeError("metadata boom") + + monkeypatch.setattr(oidc_module, "sync_request", _boom_sync_request) + with pytest.raises(ValueError) as exc_info: + oidc_module.OIDCClient._load_oidc_metadata("https://issuer.example") + assert str(exc_info.value) == "Failed to fetch OIDC metadata: metadata boom" + + +@pytest.mark.p2 +def test_parse_id_token_success_and_error(monkeypatch): + _, oidc_module = _load_auth_modules(monkeypatch) + client = _make_client(monkeypatch, oidc_module) + + monkeypatch.setattr(oidc_module.jwt, "get_unverified_header", lambda _token: {}) + + seen = {} + + class _JwkClient(_DummyJwkClient): + def __init__(self, jwks_uri): + super().__init__(jwks_uri) + seen["jwks_uri"] = jwks_uri + + def get_signing_key_from_jwt(self, id_token): + seen["id_token"] = id_token + return super().get_signing_key_from_jwt(id_token) + + monkeypatch.setattr(oidc_module.jwt, "PyJWKClient", _JwkClient) + + def _decode(id_token, key, algorithms, audience, issuer): + seen.update( + { + "decode_id_token": id_token, + "decode_key": key, + "algorithms": algorithms, + "audience": audience, + "issuer": issuer, + } + ) + return {"sub": "user-1", "email": "id@example.com"} + + monkeypatch.setattr(oidc_module.jwt, "decode", _decode) + parsed = client.parse_id_token("id-token-1") + + assert parsed["sub"] == "user-1" + assert seen["jwks_uri"] == "https://issuer.example/jwks" + assert seen["decode_key"] == "dummy-signing-key" + assert seen["algorithms"] == ["RS256"] + assert seen["audience"] == "client-1" + assert seen["issuer"] == "https://issuer.example" + + def _raise_decode(*_args, **_kwargs): + raise RuntimeError("decode boom") + + monkeypatch.setattr(oidc_module.jwt, "decode", _raise_decode) + with pytest.raises(ValueError) as exc_info: + client.parse_id_token("id-token-2") + assert str(exc_info.value) == "Error parsing ID Token: decode boom" + + +@pytest.mark.p2 +def test_fetch_user_info_merges_id_token_and_oauth_userinfo(monkeypatch): + oauth_module, oidc_module = _load_auth_modules(monkeypatch) + client = _make_client(monkeypatch, oidc_module) + + monkeypatch.setattr( + oidc_module.OIDCClient, + "parse_id_token", + lambda self, _id_token: {"picture": "id-picture", "email": "id@example.com"}, + ) + + def _fake_parent_fetch(self, access_token, **_kwargs): + assert access_token == "access-1" + return oauth_module.UserInfo( + email="oauth@example.com", + username="oauth-user", + nickname="oauth-nick", + avatar_url=None, + ) + + monkeypatch.setattr(oauth_module.OAuthClient, "fetch_user_info", _fake_parent_fetch) + + info = client.fetch_user_info("access-1", id_token="id-token") + + assert info.email == "oauth@example.com" + assert info.username == "oauth-user" + assert info.nickname == "oauth-nick" + assert info.avatar_url == "id-picture" + + +@pytest.mark.p2 +def test_async_fetch_user_info_merges_id_token_and_oauth_userinfo(monkeypatch): + oauth_module, oidc_module = _load_auth_modules(monkeypatch) + client = _make_client(monkeypatch, oidc_module) + + monkeypatch.setattr( + oidc_module.OIDCClient, + "parse_id_token", + lambda self, _id_token: {"picture": "id-picture-async", "email": "id-async@example.com"}, + ) + + async def _fake_parent_async_fetch(self, access_token, **_kwargs): + assert access_token == "access-2" + return oauth_module.UserInfo( + email="oauth-async@example.com", + username="oauth-async-user", + nickname="oauth-async-nick", + avatar_url=None, + ) + + monkeypatch.setattr(oauth_module.OAuthClient, "async_fetch_user_info", _fake_parent_async_fetch) + + info = asyncio.run(client.async_fetch_user_info("access-2", id_token="id-token")) + + assert info.email == "oauth-async@example.com" + assert info.username == "oauth-async-user" + assert info.nickname == "oauth-async-nick" + assert info.avatar_url == "id-picture-async" + + +@pytest.mark.p2 +def test_normalize_user_info_passthrough(monkeypatch): + oauth_module, oidc_module = _load_auth_modules(monkeypatch) + client = _make_client(monkeypatch, oidc_module) + + result = client.normalize_user_info( + { + "email": "user@example.com", + "username": "user", + "nickname": "User", + "picture": "picture-url", + } + ) + + assert isinstance(result, oauth_module.UserInfo) + assert result.to_dict() == { + "email": "user@example.com", + "username": "user", + "nickname": "User", + "avatar_url": "picture-url", + } + + +@pytest.mark.p2 +def test_get_auth_client_type_inference_and_unsupported(monkeypatch): + auth_module = _load_auth_init_module(monkeypatch) + + class _FakeOAuth2Client: + def __init__(self, config): + self.config = config + + class _FakeOidcClient: + def __init__(self, config): + self.config = config + + class _FakeGithubClient: + def __init__(self, config): + self.config = config + + monkeypatch.setattr( + auth_module, + "CLIENT_TYPES", + { + "oauth2": _FakeOAuth2Client, + "oidc": _FakeOidcClient, + "github": _FakeGithubClient, + }, + ) + + oidc_client = auth_module.get_auth_client({"issuer": "https://issuer.example"}) + assert isinstance(oidc_client, _FakeOidcClient) + + oauth_client = auth_module.get_auth_client({}) + assert isinstance(oauth_client, _FakeOAuth2Client) + + with pytest.raises(ValueError, match="Unsupported type: invalid"): + auth_module.get_auth_client({"type": "invalid"}) + + +@pytest.mark.p2 +def test_github_oauth_client_init_and_normalize_unit(monkeypatch): + github_module = _load_github_module(monkeypatch) + + client = github_module.GithubOAuthClient(_base_config()) + assert client.authorization_url == "https://github.com/login/oauth/authorize" + assert client.token_url == "https://github.com/login/oauth/access_token" + assert client.userinfo_url == "https://api.github.com/user" + assert client.scope == "user:email" + + normalized = client.normalize_user_info( + { + "email": "octo@example.com", + "login": "octocat", + "name": "Octo Cat", + "avatar_url": "https://avatar.example/octocat.png", + } + ) + assert normalized.to_dict() == { + "email": "octo@example.com", + "username": "octocat", + "nickname": "Octo Cat", + "avatar_url": "https://avatar.example/octocat.png", + } + + normalized_fallback = client.normalize_user_info({"email": "fallback@example.com"}) + assert normalized_fallback.to_dict() == { + "email": "fallback@example.com", + "username": "fallback", + "nickname": "fallback", + "avatar_url": "", + } + + +@pytest.mark.p2 +def test_github_fetch_user_info_sync_success_and_error_unit(monkeypatch): + github_module = _load_github_module(monkeypatch) + client = github_module.GithubOAuthClient(_base_config()) + + calls = [] + + def _fake_sync_request(method, url, headers=None, timeout=None): + calls.append((method, url, headers, timeout)) + if url.endswith("/emails"): + return _FakeResponse( + [ + {"email": "other@example.com", "primary": False}, + {"email": "octo@example.com", "primary": True}, + ] + ) + return _FakeResponse({"login": "octocat", "name": "Octo Cat", "avatar_url": "https://avatar.example/octocat.png"}) + + monkeypatch.setattr(github_module, "sync_request", _fake_sync_request) + info = client.fetch_user_info("sync-token") + + assert info.to_dict() == { + "email": "octo@example.com", + "username": "octocat", + "nickname": "Octo Cat", + "avatar_url": "https://avatar.example/octocat.png", + } + assert [call[1] for call in calls] == [ + "https://api.github.com/user", + "https://api.github.com/user/emails", + ] + assert all(call[2]["Authorization"] == "Bearer sync-token" for call in calls) + assert all(call[3] == 7 for call in calls) + + def _sync_request_raises(*_args, **_kwargs): + return _FakeResponse(err=RuntimeError("status boom")) + + monkeypatch.setattr(github_module, "sync_request", _sync_request_raises) + with pytest.raises(ValueError, match="Failed to fetch github user info: status boom"): + client.fetch_user_info("sync-token") + + +@pytest.mark.p2 +def test_github_fetch_user_info_async_success_and_error_unit(monkeypatch): + github_module = _load_github_module(monkeypatch) + client = github_module.GithubOAuthClient(_base_config()) + + calls = [] + + async def _fake_async_request(method, url, headers=None, **kwargs): + calls.append((method, url, headers, kwargs.get("timeout"))) + if url.endswith("/emails"): + return _FakeResponse( + [ + {"email": "other@example.com", "primary": False}, + {"email": "octo-async@example.com", "primary": True}, + ] + ) + return _FakeResponse( + {"login": "octocat-async", "name": "Octo Async", "avatar_url": "https://avatar.example/octo-async.png"} + ) + + monkeypatch.setattr(github_module, "async_request", _fake_async_request) + info = asyncio.run(client.async_fetch_user_info("async-token")) + + assert info.to_dict() == { + "email": "octo-async@example.com", + "username": "octocat-async", + "nickname": "Octo Async", + "avatar_url": "https://avatar.example/octo-async.png", + } + assert [call[1] for call in calls] == [ + "https://api.github.com/user", + "https://api.github.com/user/emails", + ] + assert all(call[2]["Authorization"] == "Bearer async-token" for call in calls) + assert all(call[3] == 7 for call in calls) + + async def _async_request_raises(*_args, **_kwargs): + return _FakeResponse(err=RuntimeError("async status boom")) + + monkeypatch.setattr(github_module, "async_request", _async_request_raises) + with pytest.raises(ValueError, match="Failed to fetch github user info: async status boom"): + asyncio.run(client.async_fetch_user_info("async-token")) diff --git a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py new file mode 100644 index 0000000000..1361f9df56 --- /dev/null +++ b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py @@ -0,0 +1,669 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import base64 +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _Vec(list): + def __mul__(self, scalar): + return _Vec([scalar * x for x in self]) + + __rmul__ = __mul__ + + def __add__(self, other): + return _Vec([a + b for a, b in zip(self, other)]) + + def tolist(self): + return list(self) + + +class _DummyDoc: + def __init__(self, *, doc_id="doc-1", kb_id="kb-1", name="Doc", parser_id="naive"): + self.id = doc_id + self.kb_id = kb_id + self.name = name + self.parser_id = parser_id + + def to_dict(self): + return {"id": self.id, "kb_id": self.kb_id, "name": self.name} + + +class _DummyRetCode: + SUCCESS = 0 + DATA_ERROR = 102 + EXCEPTION_ERROR = 100 + + +class _DummyParserType: + QA = "qa" + NAIVE = "naive" + + +class _DummyRetriever: + async def search(self, query, _index_name, _kb_ids, highlight=None): + class _SRes: + total = 1 + ids = ["chunk-1"] + field = { + "chunk-1": { + "content_with_weight": "chunk content", + "doc_id": "doc-1", + "docnm_kwd": "Doc", + "important_kwd": ["k1"], + "question_kwd": ["q1"], + "img_id": "img-1", + "available_int": 1, + "position_int": [], + "doc_type_kwd": "text", + } + } + highlight = {"chunk-1": " highlighted content "} + + _ = (query, highlight) + return _SRes() + + +class _DummyDocStore: + def __init__(self): + self.updated = [] + self.inserted = [] + self.deleted_inputs = [] + self.to_delete = [1] + self.chunk = { + "id": "chunk-1", + "doc_id": "doc-1", + "kb_id": "kb-1", + "content_with_weight": "chunk content", + "docnm_kwd": "Doc", + "q_2_vec": [0.1, 0.2], + "content_tks": ["a"], + "content_ltks": ["b"], + "content_sm_ltks": ["c"], + } + + def get(self, *_args, **_kwargs): + return dict(self.chunk) if self.chunk is not None else None + + def update(self, condition, payload, *_args, **_kwargs): + self.updated.append((condition, payload)) + return True + + def delete(self, condition, *_args, **_kwargs): + self.deleted_inputs.append(condition) + if not self.to_delete: + return 0 + return self.to_delete.pop(0) + + def insert(self, docs, *_args, **_kwargs): + self.inserted.extend(docs) + + +class _DummyStorage: + def __init__(self): + self.put_calls = [] + self.rm_calls = [] + + def put(self, bucket, name, binary): + self.put_calls.append((bucket, name, binary)) + + def obj_exist(self, _bucket, _name): + return True + + def rm(self, bucket, name): + self.rm_calls.append((bucket, name)) + + +class _DummyTenant: + def __init__(self, tenant_id="tenant-1"): + self.tenant_id = tenant_id + + +class _DummyLLMBundle: + def __init__(self, *_args, **_kwargs): + pass + + def encode(self, _inputs): + return [_Vec([1.0, 2.0]), _Vec([3.0, 4.0])], 9 + + +class _DummyXXHash: + def __init__(self, data): + self._data = data + + def hexdigest(self): + return f"chunk-{len(self._data)}" + + +def _run(coro): + return asyncio.run(coro) + + +def _load_chunk_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + quart_mod = ModuleType("quart") + quart_mod.request = SimpleNamespace(args={}, headers={}) + monkeypatch.setitem(sys.modules, "quart", quart_mod) + + xxhash_mod = ModuleType("xxhash") + xxhash_mod.xxh64 = lambda data: _DummyXXHash(data) + monkeypatch.setitem(sys.modules, "xxhash", xxhash_mod) + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + settings_mod = ModuleType("common.settings") + settings_mod.retriever = _DummyRetriever() + settings_mod.docStoreConn = _DummyDocStore() + settings_mod.STORAGE_IMPL = _DummyStorage() + monkeypatch.setitem(sys.modules, "common.settings", settings_mod) + common_pkg.settings = settings_mod + + constants_mod = ModuleType("common.constants") + + class _DummyLLMType: + EMBEDDING = SimpleNamespace(value="embedding") + CHAT = SimpleNamespace(value="chat") + + constants_mod.RetCode = _DummyRetCode + constants_mod.LLMType = _DummyLLMType + constants_mod.ParserType = _DummyParserType + constants_mod.PAGERANK_FLD = "pagerank_flt" + monkeypatch.setitem(sys.modules, "common.constants", constants_mod) + + string_utils_mod = ModuleType("common.string_utils") + string_utils_mod.remove_redundant_spaces = lambda text: " ".join(str(text).split()) + monkeypatch.setitem(sys.modules, "common.string_utils", string_utils_mod) + + metadata_utils_mod = ModuleType("common.metadata_utils") + metadata_utils_mod.apply_meta_data_filter = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "common.metadata_utils", metadata_utils_mod) + + misc_utils_mod = ModuleType("common.misc_utils") + + async def _thread_pool_exec(func): + return func() + + misc_utils_mod.thread_pool_exec = _thread_pool_exec + monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils_mod) + + rag_pkg = ModuleType("rag") + rag_pkg.__path__ = [] + monkeypatch.setitem(sys.modules, "rag", rag_pkg) + + rag_app_pkg = ModuleType("rag.app") + rag_app_pkg.__path__ = [] + monkeypatch.setitem(sys.modules, "rag.app", rag_app_pkg) + + rag_qa_mod = ModuleType("rag.app.qa") + rag_qa_mod.rmPrefix = lambda text: str(text).strip("Q: ").strip("A: ") + rag_qa_mod.beAdoc = lambda d, q, a, _latin: {**d, "question_kwd": [q], "content_with_weight": f"{q}\n{a}"} + monkeypatch.setitem(sys.modules, "rag.app.qa", rag_qa_mod) + + rag_tag_mod = ModuleType("rag.app.tag") + rag_tag_mod.label_question = lambda *_args, **_kwargs: [] + monkeypatch.setitem(sys.modules, "rag.app.tag", rag_tag_mod) + + rag_nlp_mod = ModuleType("rag.nlp") + rag_nlp_mod.rag_tokenizer = SimpleNamespace( + tokenize=lambda text: [str(text)], + fine_grained_tokenize=lambda toks: [f"fg:{t}" for t in toks], + is_chinese=lambda _text: False, + ) + rag_nlp_mod.search = SimpleNamespace(index_name=lambda tenant_id: f"idx-{tenant_id}") + monkeypatch.setitem(sys.modules, "rag.nlp", rag_nlp_mod) + + rag_prompts_pkg = ModuleType("rag.prompts") + rag_prompts_pkg.__path__ = [] + monkeypatch.setitem(sys.modules, "rag.prompts", rag_prompts_pkg) + + rag_generator_mod = ModuleType("rag.prompts.generator") + rag_generator_mod.cross_languages = lambda *_args, **_kwargs: [] + rag_generator_mod.keyword_extraction = lambda *_args, **_kwargs: [] + monkeypatch.setitem(sys.modules, "rag.prompts.generator", rag_generator_mod) + + apps_mod = ModuleType("api.apps") + apps_mod.__path__ = [str(repo_root / "api" / "apps")] + apps_mod.current_user = SimpleNamespace(id="user-1") + apps_mod.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + + api_utils_mod = ModuleType("api.utils.api_utils") + api_utils_mod.get_json_result = lambda data=None, message="", code=0: {"code": code, "message": message, "data": data} + api_utils_mod.get_data_error_result = lambda message="": {"code": _DummyRetCode.DATA_ERROR, "message": message, "data": False} + api_utils_mod.server_error_response = lambda exc: {"code": _DummyRetCode.EXCEPTION_ERROR, "message": repr(exc), "data": False} + api_utils_mod.validate_request = lambda *_args, **_kwargs: (lambda fn: fn) + api_utils_mod.get_request_json = lambda: _AwaitableValue({}) + monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod) + + services_pkg = ModuleType("api.db.services") + services_pkg.__path__ = [] + monkeypatch.setitem(sys.modules, "api.db.services", services_pkg) + + document_service_mod = ModuleType("api.db.services.document_service") + + class _DocumentService: + decrement_calls = [] + increment_calls = [] + + @staticmethod + def get_tenant_id(_doc_id): + return "tenant-1" + + @staticmethod + def get_by_id(doc_id): + return True, _DummyDoc(doc_id=doc_id, parser_id=_DummyParserType.NAIVE) + + @staticmethod + def get_embd_id(_doc_id): + return "embed-1" + + @staticmethod + def decrement_chunk_num(*args): + _DocumentService.decrement_calls.append(args) + + @staticmethod + def increment_chunk_num(*args): + _DocumentService.increment_calls.append(args) + + document_service_mod.DocumentService = _DocumentService + monkeypatch.setitem(sys.modules, "api.db.services.document_service", document_service_mod) + services_pkg.document_service = document_service_mod + + doc_metadata_service_mod = ModuleType("api.db.services.doc_metadata_service") + doc_metadata_service_mod.DocMetadataService = type("DocMetadataService", (), {}) + monkeypatch.setitem(sys.modules, "api.db.services.doc_metadata_service", doc_metadata_service_mod) + services_pkg.doc_metadata_service = doc_metadata_service_mod + + kb_service_mod = ModuleType("api.db.services.knowledgebase_service") + + class _KnowledgebaseService: + @staticmethod + def get_kb_ids(_tenant_id): + return ["kb-1"] + + @staticmethod + def get_by_id(_kb_id): + return True, SimpleNamespace(pagerank=0.6) + + kb_service_mod.KnowledgebaseService = _KnowledgebaseService + monkeypatch.setitem(sys.modules, "api.db.services.knowledgebase_service", kb_service_mod) + services_pkg.knowledgebase_service = kb_service_mod + + llm_service_mod = ModuleType("api.db.services.llm_service") + llm_service_mod.LLMBundle = _DummyLLMBundle + monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod) + services_pkg.llm_service = llm_service_mod + + search_service_mod = ModuleType("api.db.services.search_service") + search_service_mod.SearchService = type("SearchService", (), {}) + monkeypatch.setitem(sys.modules, "api.db.services.search_service", search_service_mod) + services_pkg.search_service = search_service_mod + + user_service_mod = ModuleType("api.db.services.user_service") + + class _UserTenantService: + @staticmethod + def query(**_kwargs): + return [_DummyTenant("tenant-1")] + + user_service_mod.UserTenantService = _UserTenantService + monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod) + services_pkg.user_service = user_service_mod + + module_name = "test_chunk_routes_unit_module" + module_path = repo_root / "api" / "apps" / "chunk_app.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _set_request_json(monkeypatch, module, payload): + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(payload)) + + +@pytest.mark.p2 +def test_list_chunk_exception_branches_unit(monkeypatch): + module = _load_chunk_module(monkeypatch) + + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "keywords": "chunk", "available_int": 0}) + res = _run(module.list_chunk()) + assert res["code"] == 0, res + assert res["data"]["total"] == 1, res + assert res["data"]["chunks"][0]["available_int"] == 1, res + + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "") + _set_request_json(monkeypatch, module, {"doc_id": "doc-1"}) + res = _run(module.list_chunk()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert res["message"] == "Tenant not found!", res + + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1") + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None)) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1"}) + res = _run(module.list_chunk()) + assert res["message"] == "Document not found!", res + + async def _raise_not_found(*_args, **_kwargs): + raise Exception("x not_found y") + + monkeypatch.setattr(module.settings.retriever, "search", _raise_not_found) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyDoc())) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1"}) + res = _run(module.list_chunk()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert res["message"] == "No chunk found!", res + + async def _raise_generic(*_args, **_kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(module.settings.retriever, "search", _raise_generic) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1"}) + res = _run(module.list_chunk()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "boom" in res["message"], res + + +@pytest.mark.p2 +def test_get_chunk_sanitize_and_exception_matrix_unit(monkeypatch): + module = _load_chunk_module(monkeypatch) + module.request = SimpleNamespace(args={"chunk_id": "chunk-1"}, headers={}) + + res = module.get() + assert res["code"] == 0, res + assert "q_2_vec" not in res["data"], res + assert "content_tks" not in res["data"], res + assert "content_ltks" not in res["data"], res + assert "content_sm_ltks" not in res["data"], res + + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: []) + res = module.get() + assert res["message"] == "Tenant not found!", res + + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [_DummyTenant("tenant-1")]) + module.settings.docStoreConn.chunk = None + res = module.get() + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "Chunk not found" in res["message"], res + + def _raise_not_found(*_args, **_kwargs): + raise Exception("NotFoundError: chunk-1") + + monkeypatch.setattr(module.settings.docStoreConn, "get", _raise_not_found) + res = module.get() + assert res["code"] == module.RetCode.DATA_ERROR, res + assert res["message"] == "Chunk not found!", res + + def _raise_generic(*_args, **_kwargs): + raise RuntimeError("get boom") + + monkeypatch.setattr(module.settings.docStoreConn, "get", _raise_generic) + res = module.get() + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "get boom" in res["message"], res + + +@pytest.mark.p2 +def test_set_chunk_bytes_qa_image_and_guard_matrix_unit(monkeypatch): + module = _load_chunk_module(monkeypatch) + + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": 1}) + with pytest.raises(TypeError, match="expected string or bytes-like object"): + _run(module.set()) + + _set_request_json( + monkeypatch, + module, + {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc", "important_kwd": "bad"}, + ) + res = _run(module.set()) + assert res["message"] == "`important_kwd` should be a list", res + + _set_request_json( + monkeypatch, + module, + {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc", "question_kwd": "bad"}, + ) + res = _run(module.set()) + assert res["message"] == "`question_kwd` should be a list", res + + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "") + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc"}) + res = _run(module.set()) + assert res["message"] == "Tenant not found!", res + + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1") + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None)) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc"}) + res = _run(module.set()) + assert res["message"] == "Document not found!", res + + monkeypatch.setattr( + module.DocumentService, + "get_by_id", + lambda _doc_id: (True, _DummyDoc(doc_id="doc-1", parser_id=module.ParserType.NAIVE)), + ) + _set_request_json( + monkeypatch, + module, + { + "doc_id": "doc-1", + "chunk_id": "chunk-1", + "content_with_weight": b"bytes-content", + "important_kwd": ["important"], + "question_kwd": ["question"], + "tag_kwd": ["tag"], + "tag_feas": [0.1], + "available_int": 0, + }, + ) + res = _run(module.set()) + assert res["code"] == 0, res + assert module.settings.docStoreConn.updated[-1][1]["content_with_weight"] == "bytes-content" + + monkeypatch.setattr( + module.DocumentService, + "get_by_id", + lambda _doc_id: (True, _DummyDoc(doc_id="doc-1", parser_id=module.ParserType.QA)), + ) + _set_request_json( + monkeypatch, + module, + { + "doc_id": "doc-1", + "chunk_id": "chunk-2", + "content_with_weight": "Q:Question\nA:Answer", + "image_base64": base64.b64encode(b"image").decode("utf-8"), + "img_id": "bucket-name", + }, + ) + res = _run(module.set()) + assert res["code"] == 0, res + assert module.settings.STORAGE_IMPL.put_calls, "image storage branch should be called" + + async def _raise_thread_pool(_func): + raise RuntimeError("set tp boom") + + monkeypatch.setattr(module, "thread_pool_exec", _raise_thread_pool) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc"}) + res = _run(module.set()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "set tp boom" in res["message"], res + + +@pytest.mark.p2 +def test_switch_chunk_success_failure_and_exception_unit(monkeypatch): + module = _load_chunk_module(monkeypatch) + + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None)) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"], "available_int": 1}) + res = _run(module.switch()) + assert res["message"] == "Document not found!", res + + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyDoc())) + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1") + monkeypatch.setattr(module.settings.docStoreConn, "update", lambda *_args, **_kwargs: False) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1", "c2"], "available_int": 0}) + res = _run(module.switch()) + assert res["message"] == "Index updating failure", res + + monkeypatch.setattr(module.settings.docStoreConn, "update", lambda *_args, **_kwargs: True) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1", "c2"], "available_int": 1}) + res = _run(module.switch()) + assert res["code"] == 0, res + assert res["data"] is True, res + + async def _raise_thread_pool(_func): + raise RuntimeError("switch tp boom") + + monkeypatch.setattr(module, "thread_pool_exec", _raise_thread_pool) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"], "available_int": 1}) + res = _run(module.switch()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "switch tp boom" in res["message"], res + + +@pytest.mark.p2 +def test_rm_chunk_delete_exception_partial_compensation_and_cleanup_unit(monkeypatch): + module = _load_chunk_module(monkeypatch) + + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None)) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"]}) + res = _run(module.rm()) + assert res["message"] == "Document not found!", res + + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyDoc())) + + def _raise_delete(*_args, **_kwargs): + raise RuntimeError("delete boom") + + monkeypatch.setattr(module.settings.docStoreConn, "delete", _raise_delete) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"]}) + res = _run(module.rm()) + assert res["message"] == "Chunk deleting failure", res + + def _delete(condition, *_args, **_kwargs): + module.settings.docStoreConn.deleted_inputs.append(condition) + if not module.settings.docStoreConn.to_delete: + return 0 + return module.settings.docStoreConn.to_delete.pop(0) + + module.settings.docStoreConn.to_delete = [0] + monkeypatch.setattr(module.settings.docStoreConn, "delete", _delete) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"]}) + res = _run(module.rm()) + assert res["message"] == "Index updating failure", res + + module.settings.docStoreConn.to_delete = [1, 2] + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1", "c2", "c3"]}) + res = _run(module.rm()) + assert res["code"] == 0, res + assert module.DocumentService.decrement_calls, "decrement_chunk_num should be called" + assert len(module.settings.STORAGE_IMPL.rm_calls) >= 1 + + module.settings.docStoreConn.to_delete = [1] + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": "c1"}) + res = _run(module.rm()) + assert res["code"] == 0, res + + async def _raise_thread_pool(_func): + raise RuntimeError("rm tp boom") + + monkeypatch.setattr(module, "thread_pool_exec", _raise_thread_pool) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"]}) + res = _run(module.rm()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "rm tp boom" in res["message"], res + + +@pytest.mark.p2 +def test_create_chunk_guards_pagerank_and_success_unit(monkeypatch): + module = _load_chunk_module(monkeypatch) + module.request = SimpleNamespace(headers={"X-Request-ID": "req-1"}, args={}) + + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk", "important_kwd": "bad"}) + res = _run(module.create()) + assert res["message"] == "`important_kwd` is required to be a list", res + + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk", "question_kwd": "bad"}) + res = _run(module.create()) + assert res["message"] == "`question_kwd` is required to be a list", res + + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None)) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk"}) + res = _run(module.create()) + assert res["message"] == "Document not found!", res + + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyDoc(doc_id="doc-1"))) + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "") + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk"}) + res = _run(module.create()) + assert res["message"] == "Tenant not found!", res + + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1") + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + _set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk"}) + res = _run(module.create()) + assert res["message"] == "Knowledgebase not found!", res + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, SimpleNamespace(pagerank=0.8))) + _set_request_json( + monkeypatch, + module, + { + "doc_id": "doc-1", + "content_with_weight": "chunk", + "important_kwd": ["i1"], + "question_kwd": ["q1"], + "tag_feas": [0.2], + }, + ) + res = _run(module.create()) + assert res["code"] == 0, res + assert module.settings.docStoreConn.inserted, "insert should be called" + inserted = module.settings.docStoreConn.inserted[-1] + assert "pagerank_flt" in inserted + assert module.DocumentService.increment_calls, "increment_chunk_num should be called" diff --git a/test/testcases/test_web_api/test_chunk_app/test_create_chunk.py b/test/testcases/test_web_api/test_chunk_app/test_create_chunk.py index 264200ad6a..0f84b583f9 100644 --- a/test/testcases/test_web_api/test_chunk_app/test_create_chunk.py +++ b/test/testcases/test_web_api/test_chunk_app/test_create_chunk.py @@ -148,6 +148,35 @@ class TestAddChunk: else: assert res["message"] == expected_message, res + @pytest.mark.p2 + def test_get_chunk_not_found(self, WebApiAuth): + res = get_chunk(WebApiAuth, {"chunk_id": "missing_chunk_id"}) + assert res["code"] != 0, res + assert "Chunk not found" in res["message"], res + + @pytest.mark.p2 + def test_create_chunk_with_tag_fields(self, WebApiAuth, add_document): + _, doc_id = add_document + res = list_chunks(WebApiAuth, {"doc_id": doc_id}) + if res["code"] == 0: + chunks_count = res["data"]["doc"]["chunk_num"] + else: + chunks_count = 0 + + payload = { + "doc_id": doc_id, + "content_with_weight": "chunk with tags", + "tag_feas": [0.1, 0.2], + "important_kwd": ["tag"], + "question_kwd": ["question"], + } + res = add_chunk(WebApiAuth, payload) + assert res["code"] == 0, res + assert res["data"]["chunk_id"], res + res = list_chunks(WebApiAuth, {"doc_id": doc_id}) + assert res["code"] == 0, res + assert res["data"]["doc"]["chunk_num"] == chunks_count + 1, res + @pytest.mark.p3 @pytest.mark.parametrize( "doc_id, expected_code, expected_message", diff --git a/test/testcases/test_web_api/test_chunk_app/test_list_chunks.py b/test/testcases/test_web_api/test_chunk_app/test_list_chunks.py index 33b795c184..1c05898473 100644 --- a/test/testcases/test_web_api/test_chunk_app/test_list_chunks.py +++ b/test/testcases/test_web_api/test_chunk_app/test_list_chunks.py @@ -17,7 +17,7 @@ import os from concurrent.futures import ThreadPoolExecutor, as_completed import pytest -from common import batch_add_chunks, list_chunks +from common import batch_add_chunks, list_chunks, update_chunk from configs import INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth @@ -88,6 +88,33 @@ class TestChunksList: else: assert res["message"] == expected_message, res + @pytest.mark.p2 + def test_available_int_filter(self, WebApiAuth, add_chunks): + _, doc_id, chunk_ids = add_chunks + chunk_id = chunk_ids[0] + + res = update_chunk( + WebApiAuth, + {"doc_id": doc_id, "chunk_id": chunk_id, "content_with_weight": "unchanged content", "available_int": 0}, + ) + assert res["code"] == 0, res + + from time import sleep + + sleep(1) + res = list_chunks(WebApiAuth, {"doc_id": doc_id, "available_int": 0}) + assert res["code"] == 0, res + assert len(res["data"]["chunks"]) >= 1, res + assert all(chunk["available_int"] == 0 for chunk in res["data"]["chunks"]), res + + # Restore the class-scoped fixture state for subsequent keyword cases. + res = update_chunk( + WebApiAuth, + {"doc_id": doc_id, "chunk_id": chunk_id, "content_with_weight": "chunk test 0", "available_int": 1}, + ) + assert res["code"] == 0, res + sleep(1) + @pytest.mark.p2 @pytest.mark.parametrize( "params, expected_page_size", diff --git a/test/testcases/test_web_api/test_chunk_app/test_rm_chunks.py b/test/testcases/test_web_api/test_chunk_app/test_rm_chunks.py index 7da5e51f95..b611fcd457 100644 --- a/test/testcases/test_web_api/test_chunk_app/test_rm_chunks.py +++ b/test/testcases/test_web_api/test_chunk_app/test_rm_chunks.py @@ -95,6 +95,30 @@ class TestChunksDeletion: assert len(res["data"]["chunks"]) == 0, res assert res["data"]["total"] == 0, res + @pytest.mark.p2 + def test_delete_scalar_chunk_id_payload(self, WebApiAuth, add_chunks_func): + _, doc_id, chunk_ids = add_chunks_func + payload = {"chunk_ids": chunk_ids[0], "doc_id": doc_id} + res = delete_chunks(WebApiAuth, payload) + assert res["code"] == 0, res + + res = list_chunks(WebApiAuth, {"doc_id": doc_id}) + assert res["code"] == 0, res + assert len(res["data"]["chunks"]) == 3, res + assert res["data"]["total"] == 3, res + + @pytest.mark.p2 + def test_delete_duplicate_ids_dedup_behavior(self, WebApiAuth, add_chunks_func): + _, doc_id, chunk_ids = add_chunks_func + payload = {"chunk_ids": [chunk_ids[0], chunk_ids[0]], "doc_id": doc_id} + res = delete_chunks(WebApiAuth, payload) + assert res["code"] == 0, res + + res = list_chunks(WebApiAuth, {"doc_id": doc_id}) + assert res["code"] == 0, res + assert len(res["data"]["chunks"]) == 3, res + assert res["data"]["total"] == 3, res + @pytest.mark.p3 def test_concurrent_deletion(self, WebApiAuth, add_document): count = 100 diff --git a/test/testcases/test_web_api/test_chunk_app/test_update_chunk.py b/test/testcases/test_web_api/test_chunk_app/test_update_chunk.py index f8715aec18..1681239c9d 100644 --- a/test/testcases/test_web_api/test_chunk_app/test_update_chunk.py +++ b/test/testcases/test_web_api/test_chunk_app/test_update_chunk.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import base64 import os from concurrent.futures import ThreadPoolExecutor, as_completed from random import randint @@ -154,6 +155,32 @@ class TestUpdateChunk: if chunk["chunk_id"] == chunk_id: assert chunk["available_int"] == payload["available_int"] + @pytest.mark.p2 + def test_update_chunk_qa_multiline_content(self, WebApiAuth, add_chunks): + _, doc_id, chunk_ids = add_chunks + payload = {"doc_id": doc_id, "chunk_id": chunk_ids[0], "content_with_weight": "Question line\nAnswer line"} + res = update_chunk(WebApiAuth, payload) + assert res["code"] == 0, res + + sleep(1) + res = list_chunks(WebApiAuth, {"doc_id": doc_id}) + assert res["code"] == 0, res + chunk = next(chunk for chunk in res["data"]["chunks"] if chunk["chunk_id"] == chunk_ids[0]) + assert chunk["content_with_weight"] == payload["content_with_weight"], res + + @pytest.mark.p2 + def test_update_chunk_with_image_payload(self, WebApiAuth, add_chunks): + _, doc_id, chunk_ids = add_chunks + payload = { + "doc_id": doc_id, + "chunk_id": chunk_ids[0], + "content_with_weight": "content with image", + "image_base64": base64.b64encode(b"img").decode("utf-8"), + "img_id": "bucket-name", + } + res = update_chunk(WebApiAuth, payload) + assert res["code"] == 0, res + @pytest.mark.p3 @pytest.mark.parametrize( "doc_id_param, expected_code, expected_message", diff --git a/test/testcases/test_web_api/test_connector_app/test_langfuse_app_unit.py b/test/testcases/test_web_api/test_connector_app/test_langfuse_app_unit.py new file mode 100644 index 0000000000..f86d157313 --- /dev/null +++ b/test/testcases/test_web_api/test_connector_app/test_langfuse_app_unit.py @@ -0,0 +1,219 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import asyncio +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _DummyAtomic: + def __enter__(self): + return self + + def __exit__(self, _exc_type, _exc, _tb): + return False + + +class _FakeApiError(Exception): + pass + + +class _FakeLangfuseClient: + def __init__(self, *, auth_result=True, auth_exc=None, project_payload=None): + self._auth_result = auth_result + self._auth_exc = auth_exc + if project_payload is None: + project_payload = {"data": [{"id": "project-id", "name": "project-name"}]} + self.api = SimpleNamespace( + projects=SimpleNamespace(get=lambda: SimpleNamespace(dict=lambda: project_payload)), + core=SimpleNamespace(api_error=SimpleNamespace(ApiError=_FakeApiError)), + ) + + def auth_check(self): + if self._auth_exc is not None: + raise self._auth_exc + return self._auth_result + + +def _run(coro): + return asyncio.run(coro) + + +def _load_langfuse_app(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + stub_apps = ModuleType("api.apps") + stub_apps.current_user = SimpleNamespace(id="tenant-1") + stub_apps.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", stub_apps) + + stub_langfuse = ModuleType("langfuse") + stub_langfuse.Langfuse = _FakeLangfuseClient + monkeypatch.setitem(sys.modules, "langfuse", stub_langfuse) + + module_path = repo_root / "api" / "apps" / "langfuse_app.py" + spec = importlib.util.spec_from_file_location("test_langfuse_app_unit", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + spec.loader.exec_module(module) + return module + + +@pytest.mark.p2 +def test_set_api_key_missing_fields_and_invalid_auth(monkeypatch): + module = _load_langfuse_app(monkeypatch) + monkeypatch.setattr(module.DB, "atomic", lambda: _DummyAtomic()) + + async def missing_fields(): + return {"secret_key": "", "public_key": "pub", "host": "http://host"} + + monkeypatch.setattr(module, "get_request_json", missing_fields) + res = _run(module.set_api_key.__wrapped__()) + assert res["code"] == 102 + assert res["message"] == "Missing required fields" + + async def invalid_auth(): + return {"secret_key": "sec", "public_key": "pub", "host": "http://host"} + + monkeypatch.setattr(module, "get_request_json", invalid_auth) + monkeypatch.setattr(module, "Langfuse", lambda **_kwargs: _FakeLangfuseClient(auth_result=False)) + res = _run(module.set_api_key.__wrapped__()) + assert res["code"] == 102 + assert res["message"] == "Invalid Langfuse keys" + + +@pytest.mark.p2 +def test_set_api_key_create_update_and_atomic_exception(monkeypatch): + module = _load_langfuse_app(monkeypatch) + monkeypatch.setattr(module.DB, "atomic", lambda: _DummyAtomic()) + monkeypatch.setattr(module, "Langfuse", lambda **_kwargs: _FakeLangfuseClient(auth_result=True)) + + async def payload(): + return {"secret_key": "sec", "public_key": "pub", "host": "http://host"} + + monkeypatch.setattr(module, "get_request_json", payload) + + calls = {"save": 0, "update": 0} + monkeypatch.setattr(module.TenantLangfuseService, "filter_by_tenant", lambda **_kwargs: None) + monkeypatch.setattr( + module.TenantLangfuseService, + "save", + lambda **_kwargs: calls.__setitem__("save", calls["save"] + 1), + ) + monkeypatch.setattr( + module.TenantLangfuseService, + "update_by_tenant", + lambda **_kwargs: calls.__setitem__("update", calls["update"] + 1), + ) + res = _run(module.set_api_key.__wrapped__()) + assert res["code"] == 0 + assert calls["save"] == 1 + + monkeypatch.setattr(module.TenantLangfuseService, "filter_by_tenant", lambda **_kwargs: {"id": "existing"}) + res = _run(module.set_api_key.__wrapped__()) + assert res["code"] == 0 + assert calls["update"] == 1 + + monkeypatch.setattr(module.TenantLangfuseService, "filter_by_tenant", lambda **_kwargs: None) + + def raise_save(**_kwargs): + raise RuntimeError("save failed") + + monkeypatch.setattr(module.TenantLangfuseService, "save", raise_save) + res = _run(module.set_api_key.__wrapped__()) + assert res["code"] == 100 + assert "save failed" in res["message"] + + +@pytest.mark.p2 +def test_get_api_key_no_record_invalid_auth_api_error_generic_error_success(monkeypatch): + module = _load_langfuse_app(monkeypatch) + + monkeypatch.setattr(module.TenantLangfuseService, "filter_by_tenant_with_info", lambda **_kwargs: None) + res = module.get_api_key.__wrapped__() + assert res["code"] == 0 + assert res["message"] == "Have not record any Langfuse keys." + + base_entry = {"secret_key": "sec", "public_key": "pub", "host": "http://host"} + monkeypatch.setattr(module.TenantLangfuseService, "filter_by_tenant_with_info", lambda **_kwargs: dict(base_entry)) + monkeypatch.setattr(module, "Langfuse", lambda **_kwargs: _FakeLangfuseClient(auth_result=False)) + res = module.get_api_key.__wrapped__() + assert res["code"] == 102 + assert res["message"] == "Invalid Langfuse keys loaded" + + monkeypatch.setattr( + module, + "Langfuse", + lambda **_kwargs: _FakeLangfuseClient(auth_exc=_FakeApiError("api exploded")), + ) + res = module.get_api_key.__wrapped__() + assert res["code"] == 0 + assert "Error from Langfuse" in res["message"] + + monkeypatch.setattr( + module, + "Langfuse", + lambda **_kwargs: _FakeLangfuseClient(auth_exc=RuntimeError("generic exploded")), + ) + res = module.get_api_key.__wrapped__() + assert res["code"] == 100 + assert "generic exploded" in res["message"] + + monkeypatch.setattr(module, "Langfuse", lambda **_kwargs: _FakeLangfuseClient(auth_result=True)) + res = module.get_api_key.__wrapped__() + assert res["code"] == 0 + assert res["data"]["project_id"] == "project-id" + assert res["data"]["project_name"] == "project-name" + + +@pytest.mark.p2 +def test_delete_api_key_no_record_success_exception(monkeypatch): + module = _load_langfuse_app(monkeypatch) + monkeypatch.setattr(module.DB, "atomic", lambda: _DummyAtomic()) + + monkeypatch.setattr(module.TenantLangfuseService, "filter_by_tenant", lambda **_kwargs: None) + res = module.delete_api_key.__wrapped__() + assert res["code"] == 0 + assert res["message"] == "Have not record any Langfuse keys." + + monkeypatch.setattr(module.TenantLangfuseService, "filter_by_tenant", lambda **_kwargs: {"id": "entry"}) + monkeypatch.setattr(module.TenantLangfuseService, "delete_model", lambda _entry: None) + res = module.delete_api_key.__wrapped__() + assert res["code"] == 0 + assert res["data"] is True + + def raise_delete(_entry): + raise RuntimeError("delete failed") + + monkeypatch.setattr(module.TenantLangfuseService, "delete_model", raise_delete) + res = module.delete_api_key.__wrapped__() + assert res["code"] == 100 + assert "delete failed" in res["message"] diff --git a/test/testcases/test_web_api/test_conversation_app/test_conversation_routes_unit.py b/test/testcases/test_web_api/test_conversation_app/test_conversation_routes_unit.py new file mode 100644 index 0000000000..294647541e --- /dev/null +++ b/test/testcases/test_web_api/test_conversation_app/test_conversation_routes_unit.py @@ -0,0 +1,584 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import sys +from copy import deepcopy +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +from anyio import Path as AsyncPath + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _DummyRequest: + def __init__(self, *, args=None, headers=None, form=None, files=None): + self.args = args or {} + self.headers = headers or {} + self.form = _AwaitableValue(form or {}) + self.files = _AwaitableValue(files or {}) + self.method = "POST" + self.content_length = 0 + + +class _DummyConversation: + def __init__(self, *, conv_id="conv-1", dialog_id="dialog-1", message=None, reference=None): + self.id = conv_id + self.dialog_id = dialog_id + self.message = message if message is not None else [] + self.reference = reference if reference is not None else [] + + def to_dict(self): + return { + "id": self.id, + "dialog_id": self.dialog_id, + "message": deepcopy(self.message), + "reference": deepcopy(self.reference), + } + + +class _DummyDialog: + def __init__(self, *, dialog_id="dialog-1", tenant_id="tenant-1", icon="avatar.png"): + self.id = dialog_id + self.tenant_id = tenant_id + self.icon = icon + self.prompt_config = {"prologue": "hello"} + self.llm_id = "" + self.llm_setting = {} + + def to_dict(self): + return { + "id": self.id, + "icon": self.icon, + "tenant_id": self.tenant_id, + "prompt_config": deepcopy(self.prompt_config), + } + + +class _DummyUploadedFile: + def __init__(self, filename): + self.filename = filename + self.saved_path = None + + async def save(self, path): + self.saved_path = path + await AsyncPath(path).write_bytes(b"audio-bytes") + + +def _run(coro): + return asyncio.run(coro) + + +def _load_conversation_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + deepdoc_pkg = ModuleType("deepdoc") + deepdoc_parser_pkg = ModuleType("deepdoc.parser") + deepdoc_parser_pkg.__path__ = [] + + class _StubPdfParser: + pass + + class _StubExcelParser: + pass + + class _StubDocxParser: + pass + + deepdoc_parser_pkg.PdfParser = _StubPdfParser + deepdoc_parser_pkg.ExcelParser = _StubExcelParser + deepdoc_parser_pkg.DocxParser = _StubDocxParser + deepdoc_pkg.parser = deepdoc_parser_pkg + monkeypatch.setitem(sys.modules, "deepdoc", deepdoc_pkg) + monkeypatch.setitem(sys.modules, "deepdoc.parser", deepdoc_parser_pkg) + + deepdoc_excel_module = ModuleType("deepdoc.parser.excel_parser") + deepdoc_excel_module.RAGFlowExcelParser = _StubExcelParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.excel_parser", deepdoc_excel_module) + + deepdoc_parser_utils = ModuleType("deepdoc.parser.utils") + deepdoc_parser_utils.get_text = lambda *_args, **_kwargs: "" + monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) + monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + + apps_mod = ModuleType("api.apps") + apps_mod.current_user = SimpleNamespace(id="user-1") + apps_mod.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + + module_name = "test_conversation_routes_unit_module" + module_path = repo_root / "api" / "apps" / "conversation_app.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _set_request_json(monkeypatch, module, payload): + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(deepcopy(payload))) + + +async def _read_sse_text(response): + chunks = [] + async for chunk in response.response: + if isinstance(chunk, bytes): + chunks.append(chunk.decode("utf-8")) + else: + chunks.append(chunk) + return "".join(chunks) + + +@pytest.mark.p2 +def test_set_conversation_update_create_and_errors(monkeypatch): + module = _load_conversation_module(monkeypatch) + + long_name = "n" * 300 + create_payload = { + "conversation_id": "conv-new", + "dialog_id": "dialog-1", + "is_new": True, + "name": long_name, + } + _set_request_json(monkeypatch, module, create_payload) + + saved = {} + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (True, _DummyDialog())) + monkeypatch.setattr(module.ConversationService, "save", lambda **kwargs: saved.update(kwargs) or True) + res = _run(module.set_conversation()) + assert res["code"] == 0 + assert len(res["data"]["name"]) == 255 + assert saved["user_id"] == "user-1" + + update_payload = { + "conversation_id": "conv-1", + "dialog_id": "dialog-1", + "is_new": False, + "name": "rename", + } + _set_request_json(monkeypatch, module, update_payload) + monkeypatch.setattr(module.ConversationService, "update_by_id", lambda *_args, **_kwargs: False) + res = _run(module.set_conversation()) + assert "Conversation not found" in res["message"] + + _set_request_json(monkeypatch, module, update_payload) + monkeypatch.setattr(module.ConversationService, "update_by_id", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (False, None)) + res = _run(module.set_conversation()) + assert "Fail to update" in res["message"] + + _set_request_json(monkeypatch, module, update_payload) + monkeypatch.setattr(module.ConversationService, "update_by_id", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (True, _DummyConversation(conv_id="conv-1"))) + res = _run(module.set_conversation()) + assert res["code"] == 0 + assert res["data"]["id"] == "conv-1" + + _set_request_json(monkeypatch, module, update_payload) + + def _raise_update(*_args, **_kwargs): + raise RuntimeError("update boom") + + monkeypatch.setattr(module.ConversationService, "update_by_id", _raise_update) + res = _run(module.set_conversation()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR + assert "update boom" in res["message"] + + missing_dialog_payload = { + "conversation_id": "conv-2", + "dialog_id": "dialog-missing", + "is_new": True, + "name": "create", + } + _set_request_json(monkeypatch, module, missing_dialog_payload) + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (False, None)) + res = _run(module.set_conversation()) + assert res["message"] == "Dialog not found" + + _set_request_json(monkeypatch, module, missing_dialog_payload) + + def _raise_dialog(_id): + raise RuntimeError("dialog boom") + + monkeypatch.setattr(module.DialogService, "get_by_id", _raise_dialog) + res = _run(module.set_conversation()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR + assert "dialog boom" in res["message"] + + +@pytest.mark.p2 +def test_get_and_getsse_authorization_and_reference_paths(monkeypatch): + module = _load_conversation_module(monkeypatch) + + conv = _DummyConversation(reference=[{"doc": "d"}, ["already-formatted"]]) + monkeypatch.setattr(module, "request", _DummyRequest(args={"conversation_id": "conv-1"})) + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (True, conv)) + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(icon="bot-avatar")]) + monkeypatch.setattr(module, "chunks_format", lambda _ref: [{"chunk": "normalized"}]) + + res = _run(module.get()) + assert res["code"] == 0 + assert res["data"]["avatar"] == "bot-avatar" + assert res["data"]["reference"][0]["chunks"] == [{"chunk": "normalized"}] + + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (False, None)) + res = _run(module.get()) + assert res["message"] == "Conversation not found!" + + monkeypatch.setattr(module, "request", _DummyRequest(args={"conversation_id": "conv-1"})) + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (True, conv)) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) + res = _run(module.get()) + assert res["code"] == module.RetCode.OPERATING_ERROR + assert "Only owner of conversation" in res["message"] + + def _raise_get(*_args, **_kwargs): + raise RuntimeError("get boom") + + monkeypatch.setattr(module.ConversationService, "get_by_id", _raise_get) + res = _run(module.get()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR + assert "get boom" in res["message"] + + monkeypatch.setattr(module, "request", _DummyRequest(headers={"Authorization": "Bearer"})) + res = module.getsse("dialog-1") + assert "Authorization is not valid" in res["message"] + + monkeypatch.setattr(module, "request", _DummyRequest(headers={"Authorization": "Bearer token-1"})) + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) + res = module.getsse("dialog-1") + assert "API key is invalid" in res["message"] + + monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace()]) + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (False, None)) + res = module.getsse("dialog-1") + assert res["message"] == "Dialog not found!" + + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (True, _DummyDialog())) + res = module.getsse("dialog-1") + assert res["code"] == 0 + assert res["data"]["avatar"] == "avatar.png" + assert "icon" not in res["data"] + + def _raise_getsse(_id): + raise RuntimeError("getsse boom") + + monkeypatch.setattr(module.DialogService, "get_by_id", _raise_getsse) + res = module.getsse("dialog-1") + assert res["code"] == module.RetCode.EXCEPTION_ERROR + assert "getsse boom" in res["message"] + + +@pytest.mark.p2 +def test_rm_and_list_conversation_guards(monkeypatch): + module = _load_conversation_module(monkeypatch) + + _set_request_json(monkeypatch, module, {"conversation_ids": ["conv-1"]}) + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (False, None)) + res = _run(module.rm()) + assert "Conversation not found" in res["message"] + + conv = _DummyConversation(conv_id="conv-1", dialog_id="dialog-1") + _set_request_json(monkeypatch, module, {"conversation_ids": ["conv-1"]}) + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (True, conv)) + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) + res = _run(module.rm()) + assert res["code"] == module.RetCode.OPERATING_ERROR + + deleted = [] + _set_request_json(monkeypatch, module, {"conversation_ids": ["conv-1"]}) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="dialog-1")]) + monkeypatch.setattr(module.ConversationService, "delete_by_id", lambda cid: deleted.append(cid) or True) + res = _run(module.rm()) + assert res["code"] == 0 + assert res["data"] is True + assert deleted == ["conv-1"] + + _set_request_json(monkeypatch, module, {"conversation_ids": ["conv-1"]}) + + def _raise_rm(*_args, **_kwargs): + raise RuntimeError("rm boom") + + monkeypatch.setattr(module.ConversationService, "get_by_id", _raise_rm) + res = _run(module.rm()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR + assert "rm boom" in res["message"] + + monkeypatch.setattr(module, "request", _DummyRequest(args={"dialog_id": "dialog-1"})) + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) + res = _run(module.list_conversation()) + assert res["code"] == module.RetCode.OPERATING_ERROR + assert "Only owner of dialog" in res["message"] + + monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="dialog-1")]) + monkeypatch.setattr(module.ConversationService, "model", SimpleNamespace(create_time="create_time")) + monkeypatch.setattr(module.ConversationService, "query", lambda **_kwargs: [_DummyConversation(conv_id="c1"), _DummyConversation(conv_id="c2")]) + res = _run(module.list_conversation()) + assert res["code"] == 0 + assert [x["id"] for x in res["data"]] == ["c1", "c2"] + + def _raise_list(**_kwargs): + raise RuntimeError("list boom") + + monkeypatch.setattr(module.ConversationService, "query", _raise_list) + res = _run(module.list_conversation()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR + assert "list boom" in res["message"] + + +@pytest.mark.p2 +def test_completion_stream_and_nonstream_branches(monkeypatch): + module = _load_conversation_module(monkeypatch) + + conv = _DummyConversation(conv_id="conv-1", dialog_id="dialog-1", reference=[]) + dia = _DummyDialog(dialog_id="dialog-1", tenant_id="tenant-1") + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (True, conv)) + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (True, dia)) + monkeypatch.setattr(module, "structure_answer", lambda _conv, ans, message_id, conv_id: {"answer": ans["answer"], "id": message_id, "conversation_id": conv_id, "reference": []}) + + updates = [] + monkeypatch.setattr(module.ConversationService, "update_by_id", lambda conv_id, payload: updates.append((conv_id, payload)) or True) + + stream_payload = { + "conversation_id": "conv-1", + "messages": [ + {"role": "system", "content": "ignored"}, + {"role": "assistant", "content": "ignored-first-assistant"}, + {"role": "user", "content": "hello", "id": "m-1"}, + ], + "stream": True, + } + + async def _stream_ok(_dia, sanitized, *_args, **_kwargs): + assert [m["role"] for m in sanitized] == ["user"] + yield {"answer": "sse-ok"} + + monkeypatch.setattr(module, "async_chat", _stream_ok) + _set_request_json(monkeypatch, module, stream_payload) + resp = _run(module.completion.__wrapped__()) + assert resp.headers["Content-Type"].startswith("text/event-stream") + sse_text = _run(_read_sse_text(resp)) + assert "sse-ok" in sse_text + assert '"data": true' in sse_text + assert updates + + async def _stream_error(_dia, _sanitized, *_args, **_kwargs): + raise RuntimeError("stream explode") + if False: + yield {"answer": "never"} + + monkeypatch.setattr(module, "async_chat", _stream_error) + _set_request_json(monkeypatch, module, stream_payload) + resp = _run(module.completion.__wrapped__()) + sse_text = _run(_read_sse_text(resp)) + assert "**ERROR**: stream explode" in sse_text + + async def _non_stream(_dia, _sanitized, **_kwargs): + yield {"answer": "plain-ok"} + + monkeypatch.setattr(module, "async_chat", _non_stream) + _set_request_json( + monkeypatch, + module, + { + "conversation_id": "conv-1", + "messages": [{"role": "user", "content": "plain", "id": "m-2"}], + "stream": False, + }, + ) + res = _run(module.completion.__wrapped__()) + assert res["code"] == 0 + assert res["data"]["answer"] == "plain-ok" + + monkeypatch.setattr(module.TenantLLMService, "get_api_key", lambda **_kwargs: False) + _set_request_json( + monkeypatch, + module, + { + "conversation_id": "conv-1", + "messages": [{"role": "user", "content": "embed", "id": "m-3"}], + "llm_id": "bad-model", + "stream": False, + }, + ) + res = _run(module.completion.__wrapped__()) + assert "Cannot use specified model bad-model" in res["message"] + + monkeypatch.setattr(module.TenantLLMService, "get_api_key", lambda **_kwargs: "api-key") + _set_request_json( + monkeypatch, + module, + { + "conversation_id": "conv-1", + "messages": [{"role": "user", "content": "embed", "id": "m-4"}], + "llm_id": "glm-4", + "temperature": 0.7, + "top_p": 0.2, + "stream": False, + }, + ) + res = _run(module.completion.__wrapped__()) + assert res["code"] == 0 + assert dia.llm_id == "glm-4" + assert dia.llm_setting == {"temperature": 0.7, "top_p": 0.2} + + _set_request_json( + monkeypatch, + module, + { + "conversation_id": "missing", + "messages": [{"role": "user", "content": "x", "id": "m-5"}], + "stream": False, + }, + ) + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (False, None)) + res = _run(module.completion.__wrapped__()) + assert res["message"] == "Conversation not found!" + + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (True, conv)) + monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (False, None)) + _set_request_json( + monkeypatch, + module, + { + "conversation_id": "conv-1", + "messages": [{"role": "user", "content": "x", "id": "m-6"}], + "stream": False, + }, + ) + res = _run(module.completion.__wrapped__()) + assert res["message"] == "Dialog not found!" + + monkeypatch.setattr(module.ConversationService, "get_by_id", lambda _id: (_ for _ in ()).throw(RuntimeError("completion boom"))) + _set_request_json( + monkeypatch, + module, + { + "conversation_id": "conv-1", + "messages": [{"role": "user", "content": "x", "id": "m-7"}], + "stream": False, + }, + ) + res = _run(module.completion.__wrapped__()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR + assert "completion boom" in res["message"] + + +@pytest.mark.p2 +def test_sequence2txt_validation_and_transcription_paths(monkeypatch): + module = _load_conversation_module(monkeypatch) + + monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={})) + res = _run(module.sequence2txt()) + assert "Missing 'file'" in res["message"] + + bad_file = _DummyUploadedFile("audio.txt") + monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": bad_file})) + res = _run(module.sequence2txt()) + assert "Unsupported audio format" in res["message"] + + wav_file = _DummyUploadedFile("audio.wav") + monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file})) + monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: []) + res = _run(module.sequence2txt()) + assert res["message"] == "Tenant not found!" + + wav_file = _DummyUploadedFile("audio.wav") + monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file})) + monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "asr_id": ""}]) + res = _run(module.sequence2txt()) + assert res["message"] == "No default ASR model is set" + + class _SyncAsr: + def transcription(self, _path): + return "transcribed text" + + def stream_transcription(self, _path): + return [] + + wav_file = _DummyUploadedFile("audio.wav") + monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file})) + monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "asr_id": "asr-model"}]) + monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _SyncAsr()) + monkeypatch.setattr(module.os, "remove", lambda _path: (_ for _ in ()).throw(RuntimeError("remove failed"))) + res = _run(module.sequence2txt()) + assert res["code"] == 0 + assert res["data"]["text"] == "transcribed text" + + class _StreamAsr: + def transcription(self, _path): + return "" + + def stream_transcription(self, _path): + yield {"event": "partial", "text": "hello"} + + wav_file = _DummyUploadedFile("audio.wav") + monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "true"}, files={"file": wav_file})) + monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _StreamAsr()) + resp = _run(module.sequence2txt()) + assert resp.headers["Content-Type"].startswith("text/event-stream") + sse_text = _run(_read_sse_text(resp)) + assert '"event": "partial"' in sse_text + + class _ErrorStreamAsr: + def transcription(self, _path): + return "" + + def stream_transcription(self, _path): + raise RuntimeError("stream asr boom") + + wav_file = _DummyUploadedFile("audio.wav") + monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "true"}, files={"file": wav_file})) + monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _ErrorStreamAsr()) + resp = _run(module.sequence2txt()) + sse_text = _run(_read_sse_text(resp)) + assert "stream asr boom" in sse_text + + +@pytest.mark.p2 +def test_tts_request_parse_entry(monkeypatch): + module = _load_conversation_module(monkeypatch) + _set_request_json(monkeypatch, module, {"text": "hello"}) + monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: []) + res = _run(module.tts()) + assert res["message"] == "Tenant not found!" diff --git a/test/testcases/test_web_api/test_document_app/conftest.py b/test/testcases/test_web_api/test_document_app/conftest.py index a34bc9be72..107e459517 100644 --- a/test/testcases/test_web_api/test_document_app/conftest.py +++ b/test/testcases/test_web_api/test_document_app/conftest.py @@ -15,10 +15,22 @@ # +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + import pytest from common import bulk_upload_documents, delete_document, list_documents +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + return decorator + + @pytest.fixture(scope="function") def add_document_func(request, WebApiAuth, add_dataset, ragflow_tmp_dir): def cleanup(): @@ -56,3 +68,49 @@ def add_documents_func(request, WebApiAuth, add_dataset_func, ragflow_tmp_dir): dataset_id = add_dataset_func return dataset_id, bulk_upload_documents(WebApiAuth, dataset_id, 3, ragflow_tmp_dir) + + +@pytest.fixture() +def document_app_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + deepdoc_pkg = ModuleType("deepdoc") + deepdoc_parser_pkg = ModuleType("deepdoc.parser") + deepdoc_parser_pkg.__path__ = [] + + class _StubPdfParser: + pass + + class _StubExcelParser: + pass + + deepdoc_parser_pkg.PdfParser = _StubPdfParser + deepdoc_pkg.parser = deepdoc_parser_pkg + monkeypatch.setitem(sys.modules, "deepdoc", deepdoc_pkg) + monkeypatch.setitem(sys.modules, "deepdoc.parser", deepdoc_parser_pkg) + deepdoc_excel_module = ModuleType("deepdoc.parser.excel_parser") + deepdoc_excel_module.RAGFlowExcelParser = _StubExcelParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.excel_parser", deepdoc_excel_module) + deepdoc_html_module = ModuleType("deepdoc.parser.html_parser") + + class _StubHtmlParser: + pass + + deepdoc_html_module.RAGFlowHtmlParser = _StubHtmlParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.html_parser", deepdoc_html_module) + monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + + stub_apps = ModuleType("api.apps") + stub_apps.current_user = SimpleNamespace(id="user-1") + stub_apps.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", stub_apps) + + module_path = repo_root / "api" / "apps" / "document_app.py" + spec = importlib.util.spec_from_file_location("test_document_app_unit", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + spec.loader.exec_module(module) + return module diff --git a/test/testcases/test_web_api/test_document_app/test_create_document.py b/test/testcases/test_web_api/test_document_app/test_create_document.py index df804487ba..8c39bdf4ec 100644 --- a/test/testcases/test_web_api/test_document_app/test_create_document.py +++ b/test/testcases/test_web_api/test_document_app/test_create_document.py @@ -13,7 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import asyncio import string +from types import SimpleNamespace from concurrent.futures import ThreadPoolExecutor, as_completed import pytest @@ -21,6 +23,7 @@ from common import create_document, list_kbs from configs import DOCUMENT_NAME_LIMIT, INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth from utils.file_utils import create_txt_file +from api.constants import FILE_NAME_LEN_LIMIT @pytest.mark.p1 @@ -90,3 +93,130 @@ class TestDocumentCreate: res = list_kbs(WebApiAuth, {"id": kb_id}) assert res["data"]["kbs"][0]["doc_num"] == count, res + + +def _run(coro): + return asyncio.run(coro) + + +@pytest.mark.p2 +class TestDocumentCreateUnit: + def test_missing_kb_id(self, document_app_module, monkeypatch): + module = document_app_module + + async def fake_request_json(): + return {"kb_id": "", "name": "doc.txt"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.create.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == 'Lack of "KB ID"' + + def test_filename_too_long(self, document_app_module, monkeypatch): + module = document_app_module + long_name = "a" * (FILE_NAME_LEN_LIMIT + 1) + + async def fake_request_json(): + return {"kb_id": "kb1", "name": long_name} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.create.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == f"File name must be {FILE_NAME_LEN_LIMIT} bytes or less." + + def test_filename_whitespace(self, document_app_module, monkeypatch): + module = document_app_module + + async def fake_request_json(): + return {"kb_id": "kb1", "name": " "} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.create.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == "File name can't be empty." + + def test_kb_not_found(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + + async def fake_request_json(): + return {"kb_id": "missing", "name": "doc.txt"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.create.__wrapped__()) + assert res["code"] == 102 + assert res["message"] == "Can't find this dataset!" + + def test_duplicate_name(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [object()]) + + async def fake_request_json(): + return {"kb_id": "kb1", "name": "doc.txt"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.create.__wrapped__()) + assert res["code"] == 102 + assert "Duplicated document name" in res["message"] + + def test_root_folder_missing(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: None) + + async def fake_request_json(): + return {"kb_id": "kb1", "name": "doc.txt"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.create.__wrapped__()) + assert res["code"] == 102 + assert res["message"] == "Cannot find the root folder." + + def test_kb_folder_missing(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "root"}) + monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: None) + + async def fake_request_json(): + return {"kb_id": "kb1", "name": "doc.txt"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.create.__wrapped__()) + assert res["code"] == 102 + assert res["message"] == "Cannot find the kb folder for this file." + + def test_success(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) + monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "root"}) + monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "folder"}) + + class _Doc: + def __init__(self, doc_id): + self.id = doc_id + + def to_json(self): + return {"id": self.id, "name": "doc.txt", "kb_id": "kb1"} + + def to_dict(self): + return {"id": self.id, "name": "doc.txt", "kb_id": "kb1"} + + monkeypatch.setattr(module.DocumentService, "insert", lambda _doc: _Doc("doc1")) + monkeypatch.setattr(module.FileService, "add_file_from_kb", lambda *_args, **_kwargs: None) + + async def fake_request_json(): + return {"kb_id": "kb1", "name": "doc.txt"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.create.__wrapped__()) + assert res["code"] == 0 + assert res["data"]["id"] == "doc1" diff --git a/test/testcases/test_web_api/test_document_app/test_document_metadata.py b/test/testcases/test_web_api/test_document_app/test_document_metadata.py index 6d0d1a3ae5..6808993e89 100644 --- a/test/testcases/test_web_api/test_document_app/test_document_metadata.py +++ b/test/testcases/test_web_api/test_document_app/test_document_metadata.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import asyncio +from types import SimpleNamespace + import pytest from common import ( document_change_status, @@ -241,3 +244,170 @@ class TestDocumentMetadataNegative: res = document_set_meta(WebApiAuth, {"doc_id": doc_id, "meta": "[]"}) assert res["code"] == 101, res assert "dictionary" in res["message"], res + + +def _run(coro): + return asyncio.run(coro) + + +@pytest.mark.p2 +class TestDocumentMetadataUnit: + def _allow_kb(self, module, monkeypatch, kb_id="kb1", tenant_id="tenant1"): + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id=tenant_id)]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: True if _kwargs.get("id") == kb_id else False) + + def test_filter_missing_kb_id(self, document_app_module, monkeypatch): + module = document_app_module + + async def fake_request_json(): + return {} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.get_filter()) + assert res["code"] == 101 + assert "KB ID" in res["message"] + + def test_filter_unauthorized(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant1")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: False) + + async def fake_request_json(): + return {"kb_id": "kb1"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.get_filter()) + assert res["code"] == 103 + + def test_filter_invalid_filters(self, document_app_module, monkeypatch): + module = document_app_module + self._allow_kb(module, monkeypatch) + + async def fake_request_json(): + return {"kb_id": "kb1", "run_status": ["INVALID"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.get_filter()) + assert res["code"] == 102 + assert "Invalid filter run status" in res["message"] + + async def fake_request_json_types(): + return {"kb_id": "kb1", "types": ["INVALID"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json_types) + res = _run(module.get_filter()) + assert res["code"] == 102 + assert "Invalid filter conditions" in res["message"] + + def test_filter_keywords_suffix(self, document_app_module, monkeypatch): + module = document_app_module + self._allow_kb(module, monkeypatch) + monkeypatch.setattr(module.DocumentService, "get_filter_by_kb_id", lambda *_args, **_kwargs: ({"run": {}}, 1)) + + async def fake_request_json(): + return {"kb_id": "kb1", "keywords": "ragflow", "suffix": ["txt"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.get_filter()) + assert res["code"] == 0 + assert "filter" in res["data"] + + def test_filter_exception(self, document_app_module, monkeypatch): + module = document_app_module + self._allow_kb(module, monkeypatch) + + def raise_error(*_args, **_kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(module.DocumentService, "get_filter_by_kb_id", raise_error) + + async def fake_request_json(): + return {"kb_id": "kb1"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.get_filter()) + assert res["code"] == 100 + + def test_infos_meta_fields(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True) + + class _Docs: + def dicts(self): + return [{"id": "doc1"}] + + monkeypatch.setattr(module.DocumentService, "get_by_ids", lambda _ids: _Docs()) + monkeypatch.setattr(module.DocMetadataService, "get_document_metadata", lambda _doc_id: {"author": "alice"}) + + async def fake_request_json(): + return {"doc_ids": ["doc1"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.doc_infos()) + assert res["code"] == 0 + assert res["data"][0]["meta_fields"]["author"] == "alice" + + def test_metadata_summary_missing_kb_id(self, document_app_module, monkeypatch): + module = document_app_module + + async def fake_request_json(): + return {"doc_ids": ["doc1"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.metadata_summary()) + assert res["code"] == 101 + + def test_metadata_summary_unauthorized(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant1")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: False) + + async def fake_request_json(): + return {"kb_id": "kb1", "doc_ids": ["doc1"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.metadata_summary()) + assert res["code"] == 103 + + def test_metadata_summary_success_and_exception(self, document_app_module, monkeypatch): + module = document_app_module + self._allow_kb(module, monkeypatch) + monkeypatch.setattr(module.DocMetadataService, "get_metadata_summary", lambda *_args, **_kwargs: {"author": {"alice": 1}}) + + async def fake_request_json(): + return {"kb_id": "kb1", "doc_ids": ["doc1"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.metadata_summary()) + assert res["code"] == 0 + assert "summary" in res["data"] + + def raise_error(*_args, **_kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(module.DocMetadataService, "get_metadata_summary", raise_error) + res = _run(module.metadata_summary()) + assert res["code"] == 100 + + def test_metadata_update_missing_kb_id(self, document_app_module, monkeypatch): + module = document_app_module + + async def fake_request_json(): + return {"doc_ids": ["doc1"], "updates": [], "deletes": []} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.metadata_update.__wrapped__()) + assert res["code"] == 101 + assert "KB ID" in res["message"] + + def test_metadata_update_success(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module.DocMetadataService, "batch_update_metadata", lambda *_args, **_kwargs: 1) + + async def fake_request_json(): + return {"kb_id": "kb1", "doc_ids": ["doc1"], "updates": [{"key": "author", "value": "alice"}], "deletes": []} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.metadata_update.__wrapped__()) + assert res["code"] == 0 + assert res["data"]["matched_docs"] == 1 diff --git a/test/testcases/test_web_api/test_document_app/test_list_documents.py b/test/testcases/test_web_api/test_document_app/test_list_documents.py index c90db5b33c..8115220efe 100644 --- a/test/testcases/test_web_api/test_document_app/test_list_documents.py +++ b/test/testcases/test_web_api/test_document_app/test_list_documents.py @@ -13,7 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import asyncio from concurrent.futures import ThreadPoolExecutor, as_completed +from types import SimpleNamespace import pytest from common import list_documents @@ -178,3 +180,214 @@ class TestDocumentsList: responses = list(as_completed(futures)) assert len(responses) == count, responses assert all(future.result()["code"] == 0 for future in futures), responses + + +def _run(coro): + return asyncio.run(coro) + + +class _DummyArgs(dict): + def get(self, key, default=None): + return super().get(key, default) + + +@pytest.mark.p2 +class TestDocumentsListUnit: + def _set_args(self, module, monkeypatch, **kwargs): + monkeypatch.setattr(module, "request", SimpleNamespace(args=_DummyArgs(kwargs))) + + def _allow_kb(self, module, monkeypatch, kb_id="kb1", tenant_id="tenant1"): + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id=tenant_id)]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: True if _kwargs.get("id") == kb_id else False) + + def test_missing_kb_id(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch) + + async def fake_request_json(): + return {} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 101 + assert res["message"] == 'Lack of "KB ID"' + + def test_unauthorized_dataset(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1") + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant1")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: False) + + async def fake_request_json(): + return {} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 103 + assert "Only owner of dataset" in res["message"] + + def test_return_empty_metadata_flags(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1") + self._allow_kb(module, monkeypatch) + monkeypatch.setattr(module.DocumentService, "get_by_kb_id", lambda *_args, **_kwargs: ([], 0)) + + async def fake_request_json(): + return {"return_empty_metadata": "true", "metadata": {"author": "alice"}} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 0 + + async def fake_request_json_empty(): + return {"metadata": {"empty_metadata": True, "author": "alice"}} + + monkeypatch.setattr(module, "get_request_json", fake_request_json_empty) + res = _run(module.list_docs()) + assert res["code"] == 0 + + def test_invalid_filters(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1") + self._allow_kb(module, monkeypatch) + + async def fake_request_json(): + return {"run_status": ["INVALID"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 102 + assert "Invalid filter run status" in res["message"] + + async def fake_request_json_types(): + return {"types": ["INVALID"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json_types) + res = _run(module.list_docs()) + assert res["code"] == 102 + assert "Invalid filter conditions" in res["message"] + + def test_invalid_metadata_types(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1") + self._allow_kb(module, monkeypatch) + + async def fake_request_json(): + return {"metadata_condition": "bad"} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 102 + assert "metadata_condition" in res["message"] + + async def fake_request_json_meta(): + return {"metadata": ["not", "object"]} + + monkeypatch.setattr(module, "get_request_json", fake_request_json_meta) + res = _run(module.list_docs()) + assert res["code"] == 102 + assert "metadata must be an object" in res["message"] + + def test_metadata_condition_empty_result(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1") + self._allow_kb(module, monkeypatch) + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda *_args, **_kwargs: {}) + monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: set()) + + async def fake_request_json(): + return {"metadata_condition": {"conditions": [{"name": "author", "comparison_operator": "is", "value": "alice"}]}} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 0 + assert res["data"]["total"] == 0 + + def test_metadata_values_intersection(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1") + self._allow_kb(module, monkeypatch) + metas = { + "author": {"alice": ["doc1", "doc2"]}, + "topic": {"rag": ["doc2"]}, + } + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda *_args, **_kwargs: metas) + + captured = {} + + def fake_get_by_kb_id(*_args, **_kwargs): + if len(_args) >= 10: + captured["doc_ids_filter"] = _args[9] + else: + captured["doc_ids_filter"] = None + return ([{"id": "doc2", "thumbnail": "", "parser_config": {}}], 1) + + monkeypatch.setattr(module.DocumentService, "get_by_kb_id", fake_get_by_kb_id) + + async def fake_request_json(): + return {"metadata": {"author": ["alice", " ", None], "topic": "rag"}} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 0 + assert captured["doc_ids_filter"] == ["doc2"] + + def test_metadata_intersection_empty(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1") + self._allow_kb(module, monkeypatch) + metas = { + "author": {"alice": ["doc1"]}, + "topic": {"rag": ["doc2"]}, + } + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda *_args, **_kwargs: metas) + + async def fake_request_json(): + return {"metadata": {"author": "alice", "topic": "rag"}} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 0 + assert res["data"]["total"] == 0 + + def test_desc_time_and_schema(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1", desc="false", create_time_from="150", create_time_to="250") + self._allow_kb(module, monkeypatch) + + docs = [ + {"id": "doc1", "thumbnail": "", "parser_config": {"metadata": {"a": 1}}, "create_time": 100}, + {"id": "doc2", "thumbnail": "", "parser_config": {"metadata": {"b": 2}}, "create_time": 200}, + ] + + def fake_get_by_kb_id(*_args, **_kwargs): + return (docs, 2) + + monkeypatch.setattr(module.DocumentService, "get_by_kb_id", fake_get_by_kb_id) + monkeypatch.setattr(module, "turn2jsonschema", lambda _meta: {"schema": True}) + + async def fake_request_json(): + return {} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 0 + assert len(res["data"]["docs"]) == 1 + assert res["data"]["docs"][0]["parser_config"]["metadata"] == {"schema": True} + + def test_exception_path(self, document_app_module, monkeypatch): + module = document_app_module + self._set_args(module, monkeypatch, kb_id="kb1") + self._allow_kb(module, monkeypatch) + + def raise_error(*_args, **_kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(module.DocumentService, "get_by_kb_id", raise_error) + + async def fake_request_json(): + return {} + + monkeypatch.setattr(module, "get_request_json", fake_request_json) + res = _run(module.list_docs()) + assert res["code"] == 100 diff --git a/test/testcases/test_web_api/test_document_app/test_upload_documents.py b/test/testcases/test_web_api/test_document_app/test_upload_documents.py index 220f53bdad..228310a35f 100644 --- a/test/testcases/test_web_api/test_document_app/test_upload_documents.py +++ b/test/testcases/test_web_api/test_document_app/test_upload_documents.py @@ -13,7 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import asyncio import string +from types import SimpleNamespace from concurrent.futures import ThreadPoolExecutor, as_completed import pytest @@ -21,6 +23,7 @@ from common import list_kbs, upload_documents from configs import DOCUMENT_NAME_LIMIT, INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth from utils.file_utils import create_txt_file +from api.constants import FILE_NAME_LEN_LIMIT @pytest.mark.p1 @@ -189,3 +192,288 @@ class TestDocumentsUpload: res = list_kbs(WebApiAuth) assert res["data"]["kbs"][0]["doc_num"] == count, res + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _coro(): + return self._value + + return _coro().__await__() + + +class _DummyFiles(dict): + def getlist(self, key): + value = self.get(key, []) + if isinstance(value, list): + return value + return [value] + + +class _DummyFile: + def __init__(self, filename): + self.filename = filename + self.closed = False + self.stream = self + + def close(self): + self.closed = True + + +class _DummyRequest: + def __init__(self, form=None, files=None): + self._form = form or {} + self._files = files or _DummyFiles() + + @property + def form(self): + return _AwaitableValue(self._form) + + @property + def files(self): + return _AwaitableValue(self._files) + + +def _run(coro): + return asyncio.run(coro) + + +@pytest.mark.p2 +class TestDocumentsUploadUnit: + def test_missing_kb_id(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": ""}, files=_DummyFiles())) + res = _run(module.upload.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == 'Lack of "KB ID"' + + def test_missing_file_part(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1"}, files=_DummyFiles())) + res = _run(module.upload.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == "No file part!" + + def test_empty_filename_closes_files(self, document_app_module, monkeypatch): + module = document_app_module + file_obj = _DummyFile("") + files = _DummyFiles({"file": [file_obj]}) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1"}, files=files)) + res = _run(module.upload.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == "No file selected!" + assert file_obj.closed is True + + def test_filename_too_long(self, document_app_module, monkeypatch): + module = document_app_module + long_name = "a" * (FILE_NAME_LEN_LIMIT + 1) + file_obj = _DummyFile(long_name) + files = _DummyFiles({"file": [file_obj]}) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1"}, files=files)) + res = _run(module.upload.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == f"File name must be {FILE_NAME_LEN_LIMIT} bytes or less." + + def test_invalid_kb_id_raises(self, document_app_module, monkeypatch): + module = document_app_module + file_obj = _DummyFile("ragflow_test.txt") + files = _DummyFiles({"file": [file_obj]}) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "missing"}, files=files)) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + with pytest.raises(LookupError): + _run(module.upload.__wrapped__()) + + def test_no_permission(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", parser_config={}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: False) + file_obj = _DummyFile("ragflow_test.txt") + files = _DummyFiles({"file": [file_obj]}) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1"}, files=files)) + res = _run(module.upload.__wrapped__()) + assert res["code"] == 109 + assert res["message"] == "No authorization." + + def test_thread_pool_errors(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", parser_config={}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) + + async def fake_thread_pool_exec(*_args, **_kwargs): + return (["unsupported type"], [("file1", "blob")]) + + monkeypatch.setattr(module, "thread_pool_exec", fake_thread_pool_exec) + file_obj = _DummyFile("ragflow_test.txt") + files = _DummyFiles({"file": [file_obj]}) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1"}, files=files)) + res = _run(module.upload.__wrapped__()) + assert res["code"] == 500 + assert "unsupported type" in res["message"] + assert res["data"] == ["file1"] + + def test_empty_upload_result(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", parser_config={}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) + + async def fake_thread_pool_exec(*_args, **_kwargs): + return (None, []) + + monkeypatch.setattr(module, "thread_pool_exec", fake_thread_pool_exec) + file_obj = _DummyFile("ragflow_test.txt") + files = _DummyFiles({"file": [file_obj]}) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1"}, files=files)) + res = _run(module.upload.__wrapped__()) + assert res["code"] == 102 + assert "file format" in res["message"] + + +@pytest.mark.p2 +class TestWebCrawlUnit: + def test_missing_kb_id(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "", "name": "doc", "url": "http://example.com"})) + res = _run(module.web_crawl.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == 'Lack of "KB ID"' + + def test_invalid_url(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1", "name": "doc", "url": "not-a-url"})) + res = _run(module.web_crawl.__wrapped__()) + assert res["code"] == 101 + assert res["message"] == "The URL format is invalid" + + def test_invalid_kb_id_raises(self, document_app_module, monkeypatch): + module = document_app_module + monkeypatch.setattr(module, "is_valid_url", lambda _url: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "missing", "name": "doc", "url": "http://example.com"})) + with pytest.raises(LookupError): + _run(module.web_crawl.__wrapped__()) + + def test_no_permission(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", parser_config={}) + monkeypatch.setattr(module, "is_valid_url", lambda _url: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: False) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1", "name": "doc", "url": "http://example.com"})) + res = _run(module.web_crawl.__wrapped__()) + assert res["code"] == 109 + assert res["message"] == "No authorization." + + def test_download_failure(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", parser_config={}) + monkeypatch.setattr(module, "is_valid_url", lambda _url: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module, "html2pdf", lambda _url: None) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1", "name": "doc", "url": "http://example.com"})) + res = _run(module.web_crawl.__wrapped__()) + assert res["code"] == 100 + assert "Download failure" in res["message"] + + def test_unsupported_type(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", parser_config={}) + monkeypatch.setattr(module, "is_valid_url", lambda _url: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4") + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"}) + monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"}) + monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"}) + monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: "bad.exe") + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1", "name": "doc", "url": "http://example.com"})) + res = _run(module.web_crawl.__wrapped__()) + assert res["code"] == 100 + assert "supported yet" in res["message"] + + @pytest.mark.parametrize( + "filename,filetype,expected_parser", + [ + ("image.png", "visual", "picture"), + ("sound.mp3", "aural", "audio"), + ("deck.pptx", "doc", "presentation"), + ("mail.eml", "doc", "email"), + ], + ) + def test_success_parser_overrides(self, document_app_module, monkeypatch, filename, filetype, expected_parser): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", parser_config={}) + captured = {} + + class _Storage: + def obj_exist(self, *_args, **_kwargs): + return False + + def put(self, *_args, **_kwargs): + captured["put"] = True + + def insert_doc(doc): + captured["doc"] = doc + + monkeypatch.setattr(module, "is_valid_url", lambda _url: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4") + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"}) + monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"}) + monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"}) + monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: filename) + monkeypatch.setattr(module, "filename_type", lambda _name: filetype) + monkeypatch.setattr(module, "thumbnail", lambda *_args, **_kwargs: "") + monkeypatch.setattr(module, "get_uuid", lambda: "doc-1") + monkeypatch.setattr(module.settings, "STORAGE_IMPL", _Storage()) + monkeypatch.setattr(module.DocumentService, "insert", insert_doc) + monkeypatch.setattr(module.FileService, "add_file_from_kb", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1", "name": "doc", "url": "http://example.com"})) + + res = _run(module.web_crawl.__wrapped__()) + assert res["code"] == 0 + assert captured["doc"]["parser_id"] == expected_parser + assert captured["put"] is True + + def test_exception_path(self, document_app_module, monkeypatch): + module = document_app_module + kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", parser_config={}) + + class _Storage: + def obj_exist(self, *_args, **_kwargs): + return False + + def put(self, *_args, **_kwargs): + return None + + def insert_doc(_doc): + raise RuntimeError("boom") + + monkeypatch.setattr(module, "is_valid_url", lambda _url: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4") + monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"}) + monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"}) + monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"}) + monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: "doc.pdf") + monkeypatch.setattr(module, "filename_type", lambda _name: "pdf") + monkeypatch.setattr(module, "thumbnail", lambda *_args, **_kwargs: "") + monkeypatch.setattr(module, "get_uuid", lambda: "doc-1") + monkeypatch.setattr(module.settings, "STORAGE_IMPL", _Storage()) + monkeypatch.setattr(module.DocumentService, "insert", insert_doc) + monkeypatch.setattr(module.FileService, "add_file_from_kb", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "request", _DummyRequest(form={"kb_id": "kb1", "name": "doc", "url": "http://example.com"})) + + res = _run(module.web_crawl.__wrapped__()) + assert res["code"] == 100 diff --git a/test/testcases/test_web_api/test_kb_app/test_kb_pipeline_tasks.py b/test/testcases/test_web_api/test_kb_app/test_kb_pipeline_tasks.py index 95841d528b..21fb0ec503 100644 --- a/test/testcases/test_web_api/test_kb_app/test_kb_pipeline_tasks.py +++ b/test/testcases/test_web_api/test_kb_app/test_kb_pipeline_tasks.py @@ -206,3 +206,25 @@ class TestKbPipelineLogs: res = kb_delete_pipeline_logs(WebApiAuth, params={"kb_id": kb_id}, payload={"log_ids": []}) assert res["code"] == 0, res assert res["data"] is True, res + + @pytest.mark.p3 + def test_list_pipeline_logs_missing_kb_id(self, WebApiAuth): + res = kb_list_pipeline_logs(WebApiAuth, params={}, payload={}) + assert res["code"] == 101, res + assert "KB ID" in res["message"], res + + @pytest.mark.p3 + def test_list_pipeline_logs_abnormal_date_filter(self, WebApiAuth, add_document): + kb_id, _ = add_document + res = kb_list_pipeline_logs( + WebApiAuth, + params={ + "kb_id": kb_id, + "desc": "false", + "create_date_from": "2025-01-01", + "create_date_to": "2025-02-01", + }, + payload={}, + ) + assert res["code"] == 102, res + assert "Create data filter is abnormal." in res["message"], res diff --git a/test/testcases/test_web_api/test_kb_app/test_kb_routes_unit.py b/test/testcases/test_web_api/test_kb_app/test_kb_routes_unit.py new file mode 100644 index 0000000000..060918be82 --- /dev/null +++ b/test/testcases/test_web_api/test_kb_app/test_kb_routes_unit.py @@ -0,0 +1,1048 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import inspect +import json +import sys +from copy import deepcopy +from datetime import datetime +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _DummyArgs(dict): + def getlist(self, key): + value = self.get(key) + if value is None: + return [] + if isinstance(value, list): + return value + return [value] + + +class _DummyKB: + def __init__(self, *, kb_id="kb-1", name="old_kb", tenant_id="tenant-1", pagerank=0): + self.id = kb_id + self.name = name + self.tenant_id = tenant_id + self.pagerank = pagerank + self.parser_config = {} + + def to_dict(self): + return { + "id": self.id, + "name": self.name, + "tenant_id": self.tenant_id, + "pagerank": self.pagerank, + "parser_config": deepcopy(self.parser_config), + } + + +class _DummyTask: + def __init__(self, task_id, progress): + self.id = task_id + self.progress = progress + + def to_dict(self): + return {"id": self.id, "progress": self.progress} + + +def _run(coro): + return asyncio.run(coro) + + +def _unwrap_route(func): + route_func = inspect.unwrap(func) + visited = set() + while getattr(route_func, "__closure__", None) and route_func not in visited: + visited.add(route_func) + nested = None + for cell in route_func.__closure__: + candidate = cell.cell_contents + if inspect.isfunction(candidate) and candidate is not route_func: + nested = inspect.unwrap(candidate) + break + if nested is None: + break + route_func = nested + return route_func + + +def _load_kb_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + deepdoc_pkg = ModuleType("deepdoc") + deepdoc_parser_pkg = ModuleType("deepdoc.parser") + deepdoc_parser_pkg.__path__ = [] + + class _StubPdfParser: + pass + + class _StubExcelParser: + pass + + class _StubDocxParser: + pass + + deepdoc_parser_pkg.PdfParser = _StubPdfParser + deepdoc_parser_pkg.ExcelParser = _StubExcelParser + deepdoc_parser_pkg.DocxParser = _StubDocxParser + deepdoc_pkg.parser = deepdoc_parser_pkg + monkeypatch.setitem(sys.modules, "deepdoc", deepdoc_pkg) + monkeypatch.setitem(sys.modules, "deepdoc.parser", deepdoc_parser_pkg) + + deepdoc_excel_module = ModuleType("deepdoc.parser.excel_parser") + deepdoc_excel_module.RAGFlowExcelParser = _StubExcelParser + monkeypatch.setitem(sys.modules, "deepdoc.parser.excel_parser", deepdoc_excel_module) + + deepdoc_parser_utils = ModuleType("deepdoc.parser.utils") + deepdoc_parser_utils.get_text = lambda *_args, **_kwargs: "" + monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) + monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + + apps_mod = ModuleType("api.apps") + apps_mod.current_user = SimpleNamespace(id="user-1") + apps_mod.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + + module_name = "test_kb_routes_unit_module" + module_path = repo_root / "api" / "apps" / "kb_app.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _set_request_json(monkeypatch, module, payload): + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(deepcopy(payload))) + + +def _set_request_args(monkeypatch, module, args): + monkeypatch.setattr(module, "request", SimpleNamespace(args=_DummyArgs(args))) + + +def _base_update_payload(**kwargs): + payload = {"kb_id": "kb-1", "name": "new_kb", "description": "", "parser_id": "naive"} + payload.update(kwargs) + return payload + + +@pytest.mark.p2 +def test_create_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + _set_request_json(monkeypatch, module, {"name": "early"}) + monkeypatch.setattr(module.KnowledgebaseService, "create_with_name", lambda **_kwargs: (False, {"code": 777, "message": "early"})) + res = _run(inspect.unwrap(module.create)()) + assert res["code"] == 777, res + + _set_request_json(monkeypatch, module, {"name": "save-fail"}) + monkeypatch.setattr(module.KnowledgebaseService, "create_with_name", lambda **_kwargs: (True, {"id": "kb-1"})) + monkeypatch.setattr(module.KnowledgebaseService, "save", lambda **_kwargs: False) + res = _run(inspect.unwrap(module.create)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + + _set_request_json(monkeypatch, module, {"name": "save-ok"}) + monkeypatch.setattr(module.KnowledgebaseService, "save", lambda **_kwargs: True) + res = _run(inspect.unwrap(module.create)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["kb_id"] == "kb-1", res + + _set_request_json(monkeypatch, module, {"name": "save-ex"}) + def _raise_save(**_kwargs): + raise RuntimeError("save boom") + monkeypatch.setattr(module.KnowledgebaseService, "save", _raise_save) + res = _run(inspect.unwrap(module.create)()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "save boom" in res["message"], res + + +@pytest.mark.p2 +def test_update_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + update_route = _unwrap_route(module.update) + + _set_request_json(monkeypatch, module, _base_update_payload(name=1)) + res = _run(update_route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "must be string" in res["message"], res + + _set_request_json(monkeypatch, module, _base_update_payload(name=" ")) + res = _run(update_route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "can't be empty" in res["message"], res + + _set_request_json(monkeypatch, module, _base_update_payload(name="a" * 129)) + res = _run(update_route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "large than" in res["message"], res + + monkeypatch.setattr(module.settings, "DOC_ENGINE_INFINITY", True) + _set_request_json(monkeypatch, module, _base_update_payload(parser_id="tag")) + res = _run(update_route()) + assert res["code"] == module.RetCode.OPERATING_ERROR, res + + _set_request_json(monkeypatch, module, _base_update_payload(pagerank=50)) + res = _run(update_route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "pagerank" in res["message"], res + + monkeypatch.setattr(module.settings, "DOC_ENGINE_INFINITY", False) + monkeypatch.setattr(module.KnowledgebaseService, "accessible4deletion", lambda *_args, **_kwargs: False) + _set_request_json(monkeypatch, module, _base_update_payload()) + res = _run(update_route()) + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible4deletion", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: []) + _set_request_json(monkeypatch, module, _base_update_payload()) + res = _run(update_route()) + assert res["code"] == module.RetCode.OPERATING_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **kwargs: [SimpleNamespace(id="kb-1")] if kwargs.get("created_by") else []) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + _set_request_json(monkeypatch, module, _base_update_payload()) + res = _run(update_route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Can't find this dataset" in res["message"], res + + kb = _DummyKB(kb_id="kb-1", name="old_name", pagerank=0) + def _query_duplicate(**kwargs): + if kwargs.get("created_by"): + return [SimpleNamespace(id="kb-1")] + if kwargs.get("name"): + return [SimpleNamespace(id="dup")] + return [] + monkeypatch.setattr(module.KnowledgebaseService, "query", _query_duplicate) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module.FileService, "filter_update", lambda *_args, **_kwargs: None) + _set_request_json(monkeypatch, module, _base_update_payload(name="new_name")) + res = _run(update_route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Duplicated dataset name" in res["message"], res + + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **kwargs: [SimpleNamespace(id="kb-1")] if kwargs.get("created_by") else []) + monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: False) + _set_request_json(monkeypatch, module, _base_update_payload(name="new_name", connectors=["c1"])) + res = _run(update_route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + + async def _thread_pool_exec(func, *args, **kwargs): + return func(*args, **kwargs) + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_exec) + monkeypatch.setattr(module.settings, "docStoreConn", SimpleNamespace(update=lambda *_args, **_kwargs: True)) + monkeypatch.setattr(module.search, "index_name", lambda _tenant: "idx") + monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.Connector2KbService, "link_connectors", lambda *_args, **_kwargs: ["warn"]) + monkeypatch.setattr(module.logging, "error", lambda *_args, **_kwargs: None) + + kb_first = _DummyKB(kb_id="kb-1", name="old_name", pagerank=0) + kb_second = _DummyKB(kb_id="kb-1", name="new_kb", pagerank=50) + get_by_id_results = [(True, kb_first), (True, kb_second)] + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: get_by_id_results.pop(0)) + _set_request_json(monkeypatch, module, _base_update_payload(name="new_kb", pagerank=50, connectors=["conn-1"])) + res = _run(update_route()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["connectors"] == ["conn-1"], res + + kb_first = _DummyKB(kb_id="kb-1", name="old_name", pagerank=50) + kb_second = _DummyKB(kb_id="kb-1", name="new_kb", pagerank=0) + get_by_id_results = [(True, kb_first), (True, kb_second)] + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: get_by_id_results.pop(0)) + monkeypatch.setattr(module.Connector2KbService, "link_connectors", lambda *_args, **_kwargs: []) + _set_request_json(monkeypatch, module, _base_update_payload(name="new_kb", pagerank=0)) + res = _run(update_route()) + assert res["code"] == module.RetCode.SUCCESS, res + + kb_first = _DummyKB(kb_id="kb-1", name="old_name", pagerank=0) + get_by_id_results = [(True, kb_first), (False, None)] + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: get_by_id_results.pop(0)) + _set_request_json(monkeypatch, module, _base_update_payload(name="new_kb")) + res = _run(update_route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Database error" in res["message"], res + + def _raise_query(**_kwargs): + raise RuntimeError("update boom") + monkeypatch.setattr(module.KnowledgebaseService, "query", _raise_query) + _set_request_json(monkeypatch, module, _base_update_payload()) + res = _run(update_route()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "update boom" in res["message"], res + + +@pytest.mark.p2 +def test_update_metadata_setting_not_found(monkeypatch): + module = _load_kb_module(monkeypatch) + _set_request_json(monkeypatch, module, {"kb_id": "missing-kb", "metadata": {}}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + res = _run(inspect.unwrap(module.update_metadata_setting)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Database error" in res["message"], res + + +@pytest.mark.p2 +def test_detail_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + _set_request_args(monkeypatch, module, {"kb_id": "kb-1"}) + monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: []) + res = inspect.unwrap(module.detail)() + assert res["code"] == module.RetCode.OPERATING_ERROR, res + + _set_request_args(monkeypatch, module, {"kb_id": "kb-1"}) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [SimpleNamespace(id="kb-1")]) + monkeypatch.setattr(module.KnowledgebaseService, "get_detail", lambda _kb_id: None) + res = inspect.unwrap(module.detail)() + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Can't find this dataset" in res["message"], res + + finish_at = datetime(2025, 1, 1, 12, 30, 0) + kb_detail = { + "id": "kb-1", + "parser_config": {"metadata": {"x": "y"}}, + "graphrag_task_finish_at": finish_at, + "raptor_task_finish_at": finish_at, + "mindmap_task_finish_at": finish_at, + } + monkeypatch.setattr(module.KnowledgebaseService, "get_detail", lambda _kb_id: deepcopy(kb_detail)) + monkeypatch.setattr(module.DocumentService, "get_total_size_by_kb_id", lambda **_kwargs: 1024) + monkeypatch.setattr(module.Connector2KbService, "list_connectors", lambda _kb_id: ["conn-1"]) + monkeypatch.setattr(module, "turn2jsonschema", lambda metadata: {"type": "object", "properties": metadata}) + res = inspect.unwrap(module.detail)() + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["size"] == 1024, res + assert res["data"]["connectors"] == ["conn-1"], res + assert isinstance(res["data"]["parser_config"]["metadata"], dict), res + assert res["data"]["graphrag_task_finish_at"] == "2025-01-01 12:30:00", res + + def _raise_tenants(**_kwargs): + raise RuntimeError("detail boom") + monkeypatch.setattr(module.UserTenantService, "query", _raise_tenants) + res = inspect.unwrap(module.detail)() + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "detail boom" in res["message"], res + + +@pytest.mark.p2 +def test_list_kbs_owner_ids_and_desc(monkeypatch): + module = _load_kb_module(monkeypatch) + + _set_request_args(monkeypatch, module, {"keywords": "", "page": "1", "page_size": "2", "parser_id": "naive", "orderby": "create_time", "desc": "false"}) + _set_request_json(monkeypatch, module, {}) + monkeypatch.setattr(module.TenantService, "get_joined_tenants_by_user_id", lambda _uid: [{"tenant_id": "tenant-1"}]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_tenant_ids", lambda *_args, **_kwargs: ([{"id": "kb-1", "tenant_id": "tenant-1"}], 1)) + res = _run(inspect.unwrap(module.list_kbs)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["total"] == 1, res + + _set_request_json(monkeypatch, module, {"owner_ids": ["tenant-1"]}) + monkeypatch.setattr( + module.KnowledgebaseService, + "get_by_tenant_ids", + lambda *_args, **_kwargs: ( + [{"id": "kb-1", "tenant_id": "tenant-1"}, {"id": "kb-2", "tenant_id": "tenant-2"}], + 2, + ), + ) + res = _run(inspect.unwrap(module.list_kbs)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["total"] == 1, res + assert all(kb["tenant_id"] == "tenant-1" for kb in res["data"]["kbs"]), res + + def _raise_kb_list(*_args, **_kwargs): + raise RuntimeError("list boom") + monkeypatch.setattr(module.KnowledgebaseService, "get_by_tenant_ids", _raise_kb_list) + res = _run(inspect.unwrap(module.list_kbs)()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "list boom" in res["message"], res + + +@pytest.mark.p2 +def test_rm_and_rm_sync_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + _set_request_json(monkeypatch, module, {"kb_id": "kb-1"}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible4deletion", lambda *_args, **_kwargs: False) + res = _run(inspect.unwrap(module.rm)()) + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible4deletion", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: []) + res = _run(inspect.unwrap(module.rm)()) + assert res["code"] == module.RetCode.OPERATING_ERROR, res + + async def _thread_pool_exec(func, *args, **kwargs): + return func(*args, **kwargs) + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_exec) + + kbs = [SimpleNamespace(id="kb-1", tenant_id="tenant-1", name="kb-1")] + monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: kbs) + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [SimpleNamespace(id="doc-1")]) + monkeypatch.setattr(module.DocumentService, "remove_document", lambda *_args, **_kwargs: False) + res = _run(inspect.unwrap(module.rm)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Document removal" in res["message"], res + + monkeypatch.setattr(module.DocumentService, "remove_document", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.File2DocumentService, "get_by_document_id", lambda _doc_id: [SimpleNamespace(file_id="file-1")]) + monkeypatch.setattr(module.FileService, "filter_delete", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module.File2DocumentService, "delete_by_document_id", lambda _doc_id: None) + + class _DocStore: + def delete(self, *_args, **_kwargs): + raise RuntimeError("drop failed") + + def delete_idx(self, *_args, **_kwargs): + return True + + monkeypatch.setattr(module.settings, "docStoreConn", _DocStore()) + monkeypatch.setattr(module.search, "index_name", lambda _tenant_id: "idx") + monkeypatch.setattr(module.KnowledgebaseService, "delete_by_id", lambda _kb_id: False) + res = _run(inspect.unwrap(module.rm)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Knowledgebase removal" in res["message"], res + + class _Storage: + def __init__(self): + self.removed = [] + + def remove_bucket(self, kb_id): + self.removed.append(kb_id) + + storage = _Storage() + monkeypatch.setattr(module.settings, "STORAGE_IMPL", storage) + + class _GoodDocStore: + def delete(self, *_args, **_kwargs): + return True + + def delete_idx(self, *_args, **_kwargs): + return True + + monkeypatch.setattr(module.settings, "docStoreConn", _GoodDocStore()) + monkeypatch.setattr(module.KnowledgebaseService, "delete_by_id", lambda _kb_id: True) + res = _run(inspect.unwrap(module.rm)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"] is True, res + assert storage.removed == ["kb-1"], storage.removed + + def _raise_rm(**_kwargs): + raise RuntimeError("rm boom") + monkeypatch.setattr(module.KnowledgebaseService, "query", _raise_rm) + res = _run(inspect.unwrap(module.rm)()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "rm boom" in res["message"], res + + +@pytest.mark.p2 +def test_tags_and_meta_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False) + res = inspect.unwrap(module.list_tags)("kb-1") + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.UserTenantService, "get_tenants_by_user_id", lambda _uid: [{"tenant_id": "tenant-1"}, {"tenant_id": "tenant-2"}]) + monkeypatch.setattr(module.settings, "retriever", SimpleNamespace(all_tags=lambda tenant_id, kb_ids: [f"{tenant_id}:{kb_ids[0]}"])) + res = inspect.unwrap(module.list_tags)("kb-1") + assert res["code"] == module.RetCode.SUCCESS, res + assert len(res["data"]) == 2, res + + _set_request_args(monkeypatch, module, {"kb_ids": "kb-1,kb-2"}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda kb_id, _uid: kb_id == "kb-1") + res = inspect.unwrap(module.list_tags_from_kbs)() + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + res = inspect.unwrap(module.list_tags_from_kbs)() + assert res["code"] == module.RetCode.SUCCESS, res + assert isinstance(res["data"], list), res + + _set_request_json(monkeypatch, module, {"tags": ["a", "b"]}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False) + res = _run(inspect.unwrap(module.rm_tags)("kb-1")) + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _DummyKB(tenant_id="tenant-1"))) + monkeypatch.setattr(module.settings, "docStoreConn", SimpleNamespace(update=lambda *_args, **_kwargs: True)) + monkeypatch.setattr(module.search, "index_name", lambda _tenant_id: "idx") + res = _run(inspect.unwrap(module.rm_tags)("kb-1")) + assert res["code"] == module.RetCode.SUCCESS, res + + _set_request_json(monkeypatch, module, {"from_tag": "a", "to_tag": "b"}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False) + res = _run(inspect.unwrap(module.rename_tags)("kb-1")) + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + res = _run(inspect.unwrap(module.rename_tags)("kb-1")) + assert res["code"] == module.RetCode.SUCCESS, res + + _set_request_args(monkeypatch, module, {"kb_ids": "kb-1,kb-2"}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda kb_id, _uid: kb_id == "kb-1") + res = inspect.unwrap(module.get_meta)() + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.DocMetadataService, "get_flatted_meta_by_kbs", lambda _kb_ids: {"source": ["a"]}) + res = inspect.unwrap(module.get_meta)() + assert res["code"] == module.RetCode.SUCCESS, res + assert "source" in res["data"], res + + _set_request_args(monkeypatch, module, {"kb_id": "kb-1"}) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False) + res = inspect.unwrap(module.get_basic_info)() + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "knowledgebase_basic_info", lambda _kb_id: {"finished": 1}) + res = inspect.unwrap(module.get_basic_info)() + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["finished"] == 1, res + + +@pytest.mark.p2 +def test_knowledge_graph_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False) + res = _run(inspect.unwrap(module.knowledge_graph)("kb-1")) + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _DummyKB(tenant_id="tenant-1"))) + monkeypatch.setattr(module.search, "index_name", lambda _tenant_id: "idx") + monkeypatch.setattr(module.settings, "docStoreConn", SimpleNamespace(index_exist=lambda *_args, **_kwargs: False)) + res = _run(inspect.unwrap(module.knowledge_graph)("kb-1")) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"] == {"graph": {}, "mind_map": {}}, res + + monkeypatch.setattr(module.settings, "docStoreConn", SimpleNamespace(index_exist=lambda *_args, **_kwargs: True)) + + class _EmptyRetriever: + async def search(self, *_args, **_kwargs): + return SimpleNamespace(ids=[], field={}) + + monkeypatch.setattr(module.settings, "retriever", _EmptyRetriever()) + res = _run(inspect.unwrap(module.knowledge_graph)("kb-1")) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"] == {"graph": {}, "mind_map": {}}, res + + graph_payload = { + "nodes": [{"id": "n2", "pagerank": 2}, {"id": "n1", "pagerank": 3}], + "edges": [ + {"source": "n1", "target": "n2", "weight": 2}, + {"source": "n1", "target": "n1", "weight": 3}, + {"source": "n1", "target": "n3", "weight": 4}, + ], + } + + class _GraphRetriever: + async def search(self, *_args, **_kwargs): + return SimpleNamespace( + ids=["bad"], + field={ + "bad": {"knowledge_graph_kwd": "graph", "content_with_weight": "{bad json"}, + }, + ) + + monkeypatch.setattr(module.settings, "retriever", _GraphRetriever()) + res = _run(inspect.unwrap(module.knowledge_graph)("kb-1")) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["graph"] == {}, res + + class _GraphRetrieverSuccess: + async def search(self, *_args, **_kwargs): + return SimpleNamespace( + ids=["good"], + field={ + "good": {"knowledge_graph_kwd": "graph", "content_with_weight": json.dumps(graph_payload)}, + }, + ) + + monkeypatch.setattr(module.settings, "retriever", _GraphRetrieverSuccess()) + res = _run(inspect.unwrap(module.knowledge_graph)("kb-1")) + assert res["code"] == module.RetCode.SUCCESS, res + assert len(res["data"]["graph"]["nodes"]) == 2, res + assert len(res["data"]["graph"]["edges"]) == 1, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False) + res = inspect.unwrap(module.delete_knowledge_graph)("kb-1") + assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True) + monkeypatch.setattr(module.settings, "docStoreConn", SimpleNamespace(delete=lambda *_args, **_kwargs: True)) + res = inspect.unwrap(module.delete_knowledge_graph)("kb-1") + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"] is True, res + + +@pytest.mark.p2 +def test_list_pipeline_logs_validation_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + _set_request_args(monkeypatch, module, {}) + _set_request_json(monkeypatch, module, {}) + res = _run(inspect.unwrap(module.list_pipeline_logs)()) + assert res["code"] == module.RetCode.ARGUMENT_ERROR, res + assert "KB ID" in res["message"], res + + _set_request_args( + monkeypatch, + module, + { + "kb_id": "kb-1", + "keywords": "k", + "page": "1", + "page_size": "10", + "orderby": "create_time", + "desc": "false", + "create_date_from": "2025-02-01", + "create_date_to": "2025-01-01", + }, + ) + _set_request_json(monkeypatch, module, {}) + monkeypatch.setattr(module.PipelineOperationLogService, "get_file_logs_by_kb_id", lambda *_args, **_kwargs: ([], 0)) + res = _run(inspect.unwrap(module.list_pipeline_logs)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["total"] == 0, res + + _set_request_args( + monkeypatch, + module, + { + "kb_id": "kb-1", + "create_date_from": "2025-01-01", + "create_date_to": "2025-02-01", + }, + ) + _set_request_json(monkeypatch, module, {}) + res = _run(inspect.unwrap(module.list_pipeline_logs)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Create data filter is abnormal." in res["message"], res + + +@pytest.mark.p2 +def test_list_pipeline_logs_filter_and_exception_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + _set_request_args( + monkeypatch, + module, + { + "kb_id": "kb-1", + "page": "1", + "page_size": "10", + "desc": "false", + "create_date_from": "2025-02-01", + "create_date_to": "2025-01-01", + }, + ) + + _set_request_json(monkeypatch, module, {"operation_status": ["BAD_STATUS"]}) + res = _run(inspect.unwrap(module.list_pipeline_logs)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "operation_status" in res["message"], res + + _set_request_json(monkeypatch, module, {"types": ["bad_type"]}) + res = _run(inspect.unwrap(module.list_pipeline_logs)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Invalid filter conditions" in res["message"], res + + def _raise_file_logs(*_args, **_kwargs): + raise RuntimeError("logs boom") + + _set_request_json(monkeypatch, module, {"suffix": [".txt"]}) + monkeypatch.setattr(module.PipelineOperationLogService, "get_file_logs_by_kb_id", _raise_file_logs) + res = _run(inspect.unwrap(module.list_pipeline_logs)()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "logs boom" in res["message"], res + + +@pytest.mark.p2 +def test_list_pipeline_dataset_logs_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + _set_request_args(monkeypatch, module, {}) + _set_request_json(monkeypatch, module, {}) + res = _run(inspect.unwrap(module.list_pipeline_dataset_logs)()) + assert res["code"] == module.RetCode.ARGUMENT_ERROR, res + assert "KB ID" in res["message"], res + + _set_request_args( + monkeypatch, + module, + { + "kb_id": "kb-1", + "desc": "false", + "create_date_from": "2025-01-01", + "create_date_to": "2025-02-01", + }, + ) + _set_request_json(monkeypatch, module, {}) + res = _run(inspect.unwrap(module.list_pipeline_dataset_logs)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Create data filter is abnormal." in res["message"], res + + _set_request_args( + monkeypatch, + module, + { + "kb_id": "kb-1", + "page": "1", + "page_size": "10", + "desc": "false", + "create_date_from": "2025-02-01", + "create_date_to": "2025-01-01", + }, + ) + _set_request_json(monkeypatch, module, {"operation_status": ["NOT_A_STATUS"]}) + res = _run(inspect.unwrap(module.list_pipeline_dataset_logs)()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "operation_status" in res["message"], res + + _set_request_args( + monkeypatch, + module, + { + "kb_id": "kb-1", + "page": "1", + "page_size": "10", + "desc": "true", + "create_date_from": "2025-02-01", + "create_date_to": "2025-01-01", + }, + ) + _set_request_json(monkeypatch, module, {"operation_status": []}) + monkeypatch.setattr( + module.PipelineOperationLogService, + "get_dataset_logs_by_kb_id", + lambda *_args, **_kwargs: ([{"id": "l1"}], 1), + ) + res = _run(inspect.unwrap(module.list_pipeline_dataset_logs)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["total"] == 1, res + assert res["data"]["logs"][0]["id"] == "l1", res + + def _raise_dataset_logs(*_args, **_kwargs): + raise RuntimeError("dataset logs boom") + + monkeypatch.setattr(module.PipelineOperationLogService, "get_dataset_logs_by_kb_id", _raise_dataset_logs) + res = _run(inspect.unwrap(module.list_pipeline_dataset_logs)()) + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "dataset logs boom" in res["message"], res + + +@pytest.mark.p2 +def test_pipeline_log_detail_and_delete_routes_branches(monkeypatch): + module = _load_kb_module(monkeypatch) + + _set_request_args(monkeypatch, module, {}) + _set_request_json(monkeypatch, module, {}) + res = _run(inspect.unwrap(module.delete_pipeline_logs)()) + assert res["code"] == module.RetCode.ARGUMENT_ERROR, res + assert "KB ID" in res["message"], res + + deleted_ids = [] + + def _delete_by_ids(log_ids): + deleted_ids.extend(log_ids) + + monkeypatch.setattr(module.PipelineOperationLogService, "delete_by_ids", _delete_by_ids) + _set_request_args(monkeypatch, module, {"kb_id": "kb-1"}) + _set_request_json(monkeypatch, module, {}) + res = _run(inspect.unwrap(module.delete_pipeline_logs)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"] is True, res + assert deleted_ids == [], deleted_ids + + _set_request_json(monkeypatch, module, {"log_ids": ["l1", "l2"]}) + res = _run(inspect.unwrap(module.delete_pipeline_logs)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert deleted_ids == ["l1", "l2"], deleted_ids + + _set_request_args(monkeypatch, module, {}) + res = inspect.unwrap(module.pipeline_log_detail)() + assert res["code"] == module.RetCode.ARGUMENT_ERROR, res + assert "Pipeline log ID" in res["message"], res + + _set_request_args(monkeypatch, module, {"log_id": "missing"}) + monkeypatch.setattr(module.PipelineOperationLogService, "get_by_id", lambda _log_id: (False, None)) + res = inspect.unwrap(module.pipeline_log_detail)() + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Invalid pipeline log ID" in res["message"], res + + class _Log: + def to_dict(self): + return {"id": "log-1", "status": "ok"} + + monkeypatch.setattr(module.PipelineOperationLogService, "get_by_id", lambda _log_id: (True, _Log())) + res = inspect.unwrap(module.pipeline_log_detail)() + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["id"] == "log-1", res + + +@pytest.mark.p2 +@pytest.mark.parametrize( + "route_name,task_attr,response_key,task_type", + [ + ("run_graphrag", "graphrag_task_id", "graphrag_task_id", "graphrag"), + ("run_raptor", "raptor_task_id", "raptor_task_id", "raptor"), + ("run_mindmap", "mindmap_task_id", "mindmap_task_id", "mindmap"), + ], +) +def test_run_pipeline_task_routes_branch_matrix(monkeypatch, route_name, task_attr, response_key, task_type): + module = _load_kb_module(monkeypatch) + route = inspect.unwrap(getattr(module, route_name)) + + def _make_kb(task_id): + payload = { + "id": "kb-1", + "tenant_id": "tenant-1", + "graphrag_task_id": "", + "raptor_task_id": "", + "mindmap_task_id": "", + } + payload[task_attr] = task_id + return SimpleNamespace(**payload) + + warnings = [] + monkeypatch.setattr(module.logging, "warning", lambda msg, *_args, **_kwargs: warnings.append(msg)) + + _set_request_json(monkeypatch, module, {"kb_id": ""}) + res = _run(route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "KB ID" in res["message"], res + + _set_request_json(monkeypatch, module, {"kb_id": "kb-1"}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + res = _run(route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Invalid Knowledgebase ID" in res["message"], res + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _make_kb("task-running"))) + monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (True, SimpleNamespace(progress=0))) + res = _run(route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "already running" in res["message"], res + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _make_kb("task-stale"))) + monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (False, None)) + monkeypatch.setattr(module.DocumentService, "get_by_kb_id", lambda **_kwargs: ([], 0)) + res = _run(route()) + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "No documents in Knowledgebase kb-1" in res["message"], res + assert warnings, "Expected warning for stale task id" + + queue_calls = {} + + def _queue_stub(**kwargs): + queue_calls.update(kwargs) + return "queued-task-id" + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _make_kb(""))) + monkeypatch.setattr( + module.DocumentService, + "get_by_kb_id", + lambda **_kwargs: ([{"id": "doc-1"}, {"id": "doc-2"}], 2), + ) + monkeypatch.setattr(module, "queue_raptor_o_graphrag_tasks", _queue_stub) + monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: False) + res = _run(route()) + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"][response_key] == "queued-task-id", res + assert queue_calls["ty"] == task_type, queue_calls + assert queue_calls["doc_ids"] == ["doc-1", "doc-2"], queue_calls + + +@pytest.mark.p2 +@pytest.mark.parametrize( + "route_name,task_attr,empty_on_missing_task,error_text", + [ + ("trace_graphrag", "graphrag_task_id", True, ""), + ("trace_raptor", "raptor_task_id", False, "RAPTOR Task Not Found or Error Occurred"), + ("trace_mindmap", "mindmap_task_id", False, "Mindmap Task Not Found or Error Occurred"), + ], +) +def test_trace_pipeline_task_routes_branch_matrix(monkeypatch, route_name, task_attr, empty_on_missing_task, error_text): + module = _load_kb_module(monkeypatch) + route = inspect.unwrap(getattr(module, route_name)) + + def _make_kb(task_id): + payload = { + "id": "kb-1", + "tenant_id": "tenant-1", + "graphrag_task_id": "", + "raptor_task_id": "", + "mindmap_task_id": "", + } + payload[task_attr] = task_id + return SimpleNamespace(**payload) + + _set_request_args(monkeypatch, module, {"kb_id": ""}) + res = route() + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "KB ID" in res["message"], res + + _set_request_args(monkeypatch, module, {"kb_id": "kb-1"}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + res = route() + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Invalid Knowledgebase ID" in res["message"], res + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _make_kb(""))) + res = route() + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"] == {}, res + + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _make_kb("task-1"))) + monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (False, None)) + res = route() + if empty_on_missing_task: + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"] == {}, res + else: + assert res["code"] == module.RetCode.DATA_ERROR, res + assert error_text in res["message"], res + + monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (True, _DummyTask("task-1", 1))) + res = route() + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"]["id"] == "task-1", res + + +@pytest.mark.p2 +def test_unbind_task_branch_matrix(monkeypatch): + module = _load_kb_module(monkeypatch) + route = inspect.unwrap(module.delete_kb_task) + + _set_request_args(monkeypatch, module, {"kb_id": ""}) + res = route() + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "KB ID" in res["message"], res + + _set_request_args(monkeypatch, module, {"kb_id": "missing", "pipeline_task_type": module.PipelineTaskType.GRAPH_RAG}) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) + res = route() + assert res["code"] == module.RetCode.SUCCESS, res + assert res["data"] is True, res + + kb = SimpleNamespace( + id="kb-1", + tenant_id="tenant-1", + graphrag_task_id="graph-task", + raptor_task_id="raptor-task", + mindmap_task_id="mindmap-task", + ) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + _set_request_args(monkeypatch, module, {"kb_id": "kb-1", "pipeline_task_type": "unknown"}) + res = route() + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Invalid task type" in res["message"], res + + cancelled = [] + deleted = [] + update_payloads = [] + monkeypatch.setattr(module.REDIS_CONN, "set", lambda key, value: cancelled.append((key, value))) + monkeypatch.setattr(module.search, "index_name", lambda _tenant_id: "idx") + monkeypatch.setattr(module.settings, "docStoreConn", SimpleNamespace(delete=lambda *args, **_kwargs: deleted.append(args))) + + def _record_update(_kb_id, payload): + update_payloads.append((_kb_id, payload)) + return True + + monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", _record_update) + + _set_request_args(monkeypatch, module, {"kb_id": "kb-1", "pipeline_task_type": module.PipelineTaskType.GRAPH_RAG}) + res = route() + assert res["code"] == module.RetCode.SUCCESS, res + + _set_request_args(monkeypatch, module, {"kb_id": "kb-1", "pipeline_task_type": module.PipelineTaskType.RAPTOR}) + res = route() + assert res["code"] == module.RetCode.SUCCESS, res + + _set_request_args(monkeypatch, module, {"kb_id": "kb-1", "pipeline_task_type": module.PipelineTaskType.MINDMAP}) + res = route() + assert res["code"] == module.RetCode.SUCCESS, res + + assert ("graph-task-cancel", "x") in cancelled, cancelled + assert ("raptor-task-cancel", "x") in cancelled, cancelled + assert ("mindmap-task-cancel", "x") in cancelled, cancelled + assert len(deleted) == 2, deleted + assert any(payload.get("graphrag_task_id") == "" for _, payload in update_payloads), update_payloads + assert any(payload.get("raptor_task_id") == "" for _, payload in update_payloads), update_payloads + assert any(payload.get("mindmap_task_id") == "" for _, payload in update_payloads), update_payloads + + class _FlakyPipelineType: + def __init__(self, target): + self.target = target + self.calls = 0 + + def __eq__(self, other): + self.calls += 1 + if self.calls == 1: + return other == self.target + return False + + _set_request_args( + monkeypatch, + module, + {"kb_id": "kb-1", "pipeline_task_type": _FlakyPipelineType(module.PipelineTaskType.GRAPH_RAG)}, + ) + res = route() + assert res["code"] == module.RetCode.DATA_ERROR, res + assert "Internal Error: Invalid task type" in res["message"], res + + monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: False) + monkeypatch.setattr(module, "server_error_response", lambda e: module.get_json_result(code=module.RetCode.EXCEPTION_ERROR, message=str(e))) + _set_request_args(monkeypatch, module, {"kb_id": "kb-1", "pipeline_task_type": module.PipelineTaskType.GRAPH_RAG}) + res = route() + assert res["code"] == module.RetCode.EXCEPTION_ERROR, res + assert "cannot delete task" in res["message"], res diff --git a/test/testcases/test_web_api/test_kb_app/test_kb_tags_meta.py b/test/testcases/test_web_api/test_kb_app/test_kb_tags_meta.py index 479799ad1d..810b636de6 100644 --- a/test/testcases/test_web_api/test_kb_app/test_kb_tags_meta.py +++ b/test/testcases/test_web_api/test_kb_app/test_kb_tags_meta.py @@ -17,9 +17,11 @@ import uuid import pytest from common import ( + delete_knowledge_graph, kb_basic_info, kb_get_meta, kb_update_metadata_setting, + knowledge_graph, list_tags, list_tags_from_kbs, rename_tags, @@ -121,6 +123,20 @@ class TestAuthorization: assert res["code"] == expected_code, res assert expected_fragment in res["message"], res + @pytest.mark.p2 + @pytest.mark.parametrize("invalid_auth, expected_code, expected_fragment", INVALID_AUTH_CASES) + def test_knowledge_graph_auth_invalid(self, invalid_auth, expected_code, expected_fragment): + res = knowledge_graph(invalid_auth, "kb_id") + assert res["code"] == expected_code, res + assert expected_fragment in res["message"], res + + @pytest.mark.p2 + @pytest.mark.parametrize("invalid_auth, expected_code, expected_fragment", INVALID_AUTH_CASES) + def test_delete_knowledge_graph_auth_invalid(self, invalid_auth, expected_code, expected_fragment): + res = delete_knowledge_graph(invalid_auth, "kb_id") + assert res["code"] == expected_code, res + assert expected_fragment in res["message"], res + class TestKbTagsMeta: @pytest.mark.p2 @@ -205,6 +221,22 @@ class TestKbTagsMeta: assert res["data"]["id"] == kb_id, res assert res["data"]["parser_config"]["metadata"] == metadata, res + @pytest.mark.p2 + def test_knowledge_graph(self, WebApiAuth, add_dataset): + kb_id = add_dataset + res = knowledge_graph(WebApiAuth, kb_id) + assert res["code"] == 0, res + assert isinstance(res["data"], dict), res + assert "graph" in res["data"], res + assert "mind_map" in res["data"], res + + @pytest.mark.p2 + def test_delete_knowledge_graph(self, WebApiAuth, add_dataset): + kb_id = add_dataset + res = delete_knowledge_graph(WebApiAuth, kb_id) + assert res["code"] == 0, res + assert res["data"] is True, res + class TestKbTagsMetaNegative: @pytest.mark.p3 @@ -249,3 +281,15 @@ class TestKbTagsMetaNegative: assert res["code"] == 101, res assert "required argument are missing" in res["message"], res assert "metadata" in res["message"], res + + @pytest.mark.p3 + def test_knowledge_graph_invalid_kb(self, WebApiAuth): + res = knowledge_graph(WebApiAuth, "invalid_kb_id") + assert res["code"] == 109, res + assert "No authorization" in res["message"], res + + @pytest.mark.p3 + def test_delete_knowledge_graph_invalid_kb(self, WebApiAuth): + res = delete_knowledge_graph(WebApiAuth, "invalid_kb_id") + assert res["code"] == 109, res + assert "No authorization" in res["message"], res diff --git a/test/testcases/test_web_api/test_kb_app/test_list_kbs.py b/test/testcases/test_web_api/test_kb_app/test_list_kbs.py index 6272ea3046..530686788b 100644 --- a/test/testcases/test_web_api/test_kb_app/test_list_kbs.py +++ b/test/testcases/test_web_api/test_kb_app/test_list_kbs.py @@ -182,3 +182,20 @@ class TestDatasetsList: res = list_kbs(WebApiAuth, params) assert res["code"] == 0, res assert len(res["data"]["kbs"]) == expected_page_size, res + + @pytest.mark.p2 + def test_owner_ids_payload_mode(self, WebApiAuth): + base_res = list_kbs(WebApiAuth, {"page_size": 10}) + assert base_res["code"] == 0, base_res + assert base_res["data"]["kbs"], base_res + owner_id = base_res["data"]["kbs"][0]["tenant_id"] + + res = list_kbs( + WebApiAuth, + params={"page": 1, "page_size": 2, "desc": "false"}, + payload={"owner_ids": [owner_id]}, + ) + assert res["code"] == 0, res + assert res["data"]["total"] >= len(res["data"]["kbs"]), res + assert len(res["data"]["kbs"]) <= 2, res + assert all(kb["tenant_id"] == owner_id for kb in res["data"]["kbs"]), res diff --git a/test/testcases/test_web_api/test_llm_app/test_llm_list_unit.py b/test/testcases/test_web_api/test_llm_app/test_llm_list_unit.py new file mode 100644 index 0000000000..f0a285e016 --- /dev/null +++ b/test/testcases/test_web_api/test_llm_app/test_llm_list_unit.py @@ -0,0 +1,290 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _ExprField: + def __init__(self, name): + self.name = name + + def __eq__(self, other): + return (self.name, other) + + +class _DummyTenantLLMModel: + tenant_id = _ExprField("tenant_id") + llm_factory = _ExprField("llm_factory") + + +class _TenantLLMRow: + def __init__(self, *, llm_name, llm_factory, model_type, api_key="key", status="1"): + self.llm_name = llm_name + self.llm_factory = llm_factory + self.model_type = model_type + self.api_key = api_key + self.status = status + + def to_dict(self): + return { + "llm_name": self.llm_name, + "llm_factory": self.llm_factory, + "model_type": self.model_type, + "status": self.status, + } + + +class _LLMRow: + def __init__(self, *, llm_name, fid, model_type, status="1"): + self.llm_name = llm_name + self.fid = fid + self.model_type = model_type + self.status = status + + def to_dict(self): + return { + "llm_name": self.llm_name, + "fid": self.fid, + "model_type": self.model_type, + "status": self.status, + } + + +def _run(coro): + return asyncio.run(coro) + + +def _load_llm_app(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + quart_mod = ModuleType("quart") + quart_mod.request = SimpleNamespace(args={}) + monkeypatch.setitem(sys.modules, "quart", quart_mod) + + apps_mod = ModuleType("api.apps") + apps_mod.__path__ = [str(repo_root / "api" / "apps")] + apps_mod.login_required = lambda fn: fn + apps_mod.current_user = SimpleNamespace(id="tenant-1") + monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + + tenant_llm_mod = ModuleType("api.db.services.tenant_llm_service") + + class _StubLLMFactoriesService: + @staticmethod + def query(**_kwargs): + return [] + + class _StubTenantLLMService: + @staticmethod + def ensure_mineru_from_env(_tenant_id): + return None + + @staticmethod + def query(**_kwargs): + return [] + + @staticmethod + def get_my_llms(_tenant_id): + return [] + + @staticmethod + def save(**_kwargs): + return True + + @staticmethod + def filter_delete(_filters): + return True + + tenant_llm_mod.LLMFactoriesService = _StubLLMFactoriesService + tenant_llm_mod.TenantLLMService = _StubTenantLLMService + monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_mod) + + llm_service_mod = ModuleType("api.db.services.llm_service") + + class _StubLLMService: + @staticmethod + def get_all(): + return [] + + @staticmethod + def query(**_kwargs): + return [] + + llm_service_mod.LLMService = _StubLLMService + monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod) + + api_utils_mod = ModuleType("api.utils.api_utils") + api_utils_mod.get_allowed_llm_factories = lambda: [] + api_utils_mod.get_data_error_result = lambda message="", code=400, data=None: { + "code": code, + "message": message, + "data": data, + } + api_utils_mod.get_json_result = lambda data=None, message="", code=0: { + "code": code, + "message": message, + "data": data, + } + + async def _get_request_json(): + return {} + + api_utils_mod.get_request_json = _get_request_json + api_utils_mod.server_error_response = lambda exc: {"code": 500, "message": str(exc), "data": None} + api_utils_mod.validate_request = lambda *_args, **_kwargs: (lambda fn: fn) + monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod) + + constants_mod = ModuleType("common.constants") + constants_mod.StatusEnum = SimpleNamespace(VALID=SimpleNamespace(value="1"), INVALID=SimpleNamespace(value="0")) + constants_mod.LLMType = SimpleNamespace( + CHAT="chat", + EMBEDDING="embedding", + SPEECH2TEXT="speech2text", + IMAGE2TEXT="image2text", + RERANK="rerank", + TTS="tts", + OCR="ocr", + ) + monkeypatch.setitem(sys.modules, "common.constants", constants_mod) + + db_models_mod = ModuleType("api.db.db_models") + db_models_mod.TenantLLM = _DummyTenantLLMModel + monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod) + + base64_mod = ModuleType("rag.utils.base64_image") + base64_mod.test_image = lambda _s: _s + monkeypatch.setitem(sys.modules, "rag.utils.base64_image", base64_mod) + + rag_llm_mod = ModuleType("rag.llm") + rag_llm_mod.EmbeddingModel = {} + rag_llm_mod.ChatModel = {} + rag_llm_mod.RerankModel = {} + rag_llm_mod.CvModel = {} + rag_llm_mod.TTSModel = {} + rag_llm_mod.OcrModel = {} + rag_llm_mod.Seq2txtModel = {} + monkeypatch.setitem(sys.modules, "rag.llm", rag_llm_mod) + + module_path = repo_root / "api" / "apps" / "llm_app.py" + spec = importlib.util.spec_from_file_location("test_llm_list_unit_module", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + spec.loader.exec_module(module) + return module + + +@pytest.mark.p2 +def test_list_app_grouping_availability_and_merge(monkeypatch): + module = _load_llm_app(monkeypatch) + + ensure_calls = [] + monkeypatch.setattr(module.TenantLLMService, "ensure_mineru_from_env", lambda tenant_id: ensure_calls.append(tenant_id)) + + tenant_rows = [ + _TenantLLMRow(llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"), + _TenantLLMRow(llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"), + ] + monkeypatch.setattr(module.TenantLLMService, "query", lambda **_kwargs: tenant_rows) + + all_llms = [ + _LLMRow(llm_name="tei-embed", fid="Builtin", model_type="embedding", status="1"), + _LLMRow(llm_name="fast-emb", fid="FastEmbed", model_type="embedding", status="1"), + _LLMRow(llm_name="not-in-status", fid="Other", model_type="chat", status="1"), + ] + monkeypatch.setattr(module.LLMService, "get_all", lambda: all_llms) + + monkeypatch.setattr(module, "request", SimpleNamespace(args={})) + monkeypatch.setenv("COMPOSE_PROFILES", "tei-cpu") + monkeypatch.setenv("TEI_MODEL", "tei-embed") + + res = _run(module.list_app()) + assert res["code"] == 0 + assert ensure_calls == ["tenant-1"] + + data = res["data"] + assert {"Builtin", "FastEmbed", "CustomFactory"}.issubset(set(data.keys())) + + builtin = data["Builtin"][0] + assert builtin["llm_name"] == "tei-embed" + assert builtin["available"] is True + + fastembed = data["FastEmbed"][0] + assert fastembed["llm_name"] == "fast-emb" + assert fastembed["available"] is True + + tenant_only = data["CustomFactory"][0] + assert tenant_only["llm_name"] == "tenant-only" + assert tenant_only["available"] is True + + +@pytest.mark.p2 +def test_list_app_model_type_filter(monkeypatch): + module = _load_llm_app(monkeypatch) + + monkeypatch.setattr(module.TenantLLMService, "ensure_mineru_from_env", lambda _tenant_id: None) + monkeypatch.setattr( + module.TenantLLMService, + "query", + lambda **_kwargs: [ + _TenantLLMRow(llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"), + _TenantLLMRow(llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"), + ], + ) + monkeypatch.setattr( + module.LLMService, + "get_all", + lambda: [ + _LLMRow(llm_name="tei-embed", fid="Builtin", model_type="embedding", status="1"), + _LLMRow(llm_name="fast-emb", fid="FastEmbed", model_type="embedding", status="1"), + ], + ) + + monkeypatch.setattr(module, "request", SimpleNamespace(args={"model_type": "chat"})) + res = _run(module.list_app()) + assert res["code"] == 0 + assert list(res["data"].keys()) == ["CustomFactory"] + assert res["data"]["CustomFactory"][0]["model_type"] == "chat" + + +@pytest.mark.p2 +def test_list_app_exception_path(monkeypatch): + module = _load_llm_app(monkeypatch) + + monkeypatch.setattr(module, "request", SimpleNamespace(args={})) + monkeypatch.setattr(module.TenantLLMService, "ensure_mineru_from_env", lambda _tenant_id: None) + monkeypatch.setattr( + module.TenantLLMService, + "query", + lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("query boom")), + ) + + res = _run(module.list_app()) + assert res["code"] == 500 + assert "query boom" in res["message"] diff --git a/test/testcases/test_web_api/test_mcp_server_app/test_mcp_server_app_unit.py b/test/testcases/test_web_api/test_mcp_server_app/test_mcp_server_app_unit.py new file mode 100644 index 0000000000..49af8223af --- /dev/null +++ b/test/testcases/test_web_api/test_mcp_server_app/test_mcp_server_app_unit.py @@ -0,0 +1,708 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import asyncio +import importlib.util +import inspect +import sys +from functools import wraps +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _Field: + def __init__(self, name): + self.name = name + + def __eq__(self, other): + return (self.name, other) + + +class _DummyMCPServer: + id = _Field("id") + tenant_id = _Field("tenant_id") + + def __init__(self, **kwargs): + self.id = kwargs.get("id", "") + self.name = kwargs.get("name", "") + self.url = kwargs.get("url", "") + self.server_type = kwargs.get("server_type", "sse") + self.tenant_id = kwargs.get("tenant_id", "tenant_1") + self.variables = kwargs.get("variables", {}) + self.headers = kwargs.get("headers", {}) + + def to_dict(self): + return { + "id": self.id, + "name": self.name, + "url": self.url, + "server_type": self.server_type, + "tenant_id": self.tenant_id, + "variables": self.variables, + "headers": self.headers, + } + + +class _DummyMCPServerService: + @staticmethod + def get_servers(*_args, **_kwargs): + return [] + + @staticmethod + def get_or_none(*_args, **_kwargs): + return None + + @staticmethod + def get_by_id(*_args, **_kwargs): + return False, None + + @staticmethod + def get_by_name_and_tenant(*_args, **_kwargs): + return False, None + + @staticmethod + def insert(**_kwargs): + return True + + @staticmethod + def filter_update(*_args, **_kwargs): + return True + + @staticmethod + def delete_by_ids(*_args, **_kwargs): + return True + + +class _DummyTenantService: + @staticmethod + def get_by_id(*_args, **_kwargs): + return True, SimpleNamespace(id="tenant_1") + + +class _DummyTool: + def __init__(self, name): + self._name = name + + def model_dump(self): + return {"name": self._name} + + +class _DummyMCPToolCallSession: + def __init__(self, _mcp_server, _variables): + self._tools = [_DummyTool("tool_a"), _DummyTool("tool_b")] + + def get_tools(self, _timeout): + return self._tools + + def tool_call(self, _name, _arguments, _timeout): + return "ok" + + +def _run(coro): + return asyncio.run(coro) + + +def _set_request_json(monkeypatch, module, payload): + async def _request_json(): + return payload + + monkeypatch.setattr(module, "get_request_json", _request_json) + + +def _load_mcp_server_app(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + apps_mod = ModuleType("api.apps") + apps_mod.current_user = SimpleNamespace(id="tenant_1") + apps_mod.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + + db_models_mod = ModuleType("api.db.db_models") + db_models_mod.MCPServer = _DummyMCPServer + monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod) + + mcp_service_mod = ModuleType("api.db.services.mcp_server_service") + mcp_service_mod.MCPServerService = _DummyMCPServerService + monkeypatch.setitem(sys.modules, "api.db.services.mcp_server_service", mcp_service_mod) + + user_service_mod = ModuleType("api.db.services.user_service") + user_service_mod.TenantService = _DummyTenantService + monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod) + + mcp_conn_mod = ModuleType("common.mcp_tool_call_conn") + mcp_conn_mod.MCPToolCallSession = _DummyMCPToolCallSession + mcp_conn_mod.close_multiple_mcp_toolcall_sessions = lambda _sessions: None + monkeypatch.setitem(sys.modules, "common.mcp_tool_call_conn", mcp_conn_mod) + + api_utils_mod = ModuleType("api.utils.api_utils") + + async def _default_request_json(): + return {} + + def _get_json_result(code=0, message="success", data=None): + return {"code": code, "message": message, "data": data} + + def _get_data_error_result(code=102, message="Sorry! Data missing!"): + return {"code": code, "message": message} + + def _server_error_response(error): + return {"code": 100, "message": repr(error)} + + async def _get_mcp_tools(*_args, **_kwargs): + return {} + + def _validate_request(*_args, **_kwargs): + def _decorator(func): + @wraps(func) + async def _wrapped(*func_args, **func_kwargs): + if inspect.iscoroutinefunction(func): + return await func(*func_args, **func_kwargs) + return func(*func_args, **func_kwargs) + + return _wrapped + + return _decorator + + api_utils_mod.get_request_json = _default_request_json + api_utils_mod.get_json_result = _get_json_result + api_utils_mod.get_data_error_result = _get_data_error_result + api_utils_mod.server_error_response = _server_error_response + api_utils_mod.validate_request = _validate_request + api_utils_mod.get_mcp_tools = _get_mcp_tools + monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod) + + module_name = "test_mcp_server_app_unit_module" + module_path = repo_root / "api" / "apps" / "mcp_server_app.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +@pytest.mark.p2 +def test_list_mcp_desc_pagination_and_exception(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + monkeypatch.setattr( + module, + "request", + SimpleNamespace(args={"keywords": "k", "page": "2", "page_size": "1", "orderby": "create_time", "desc": "false"}), + ) + _set_request_json(monkeypatch, module, {"mcp_ids": []}) + monkeypatch.setattr(module.MCPServerService, "get_servers", lambda *_args, **_kwargs: [{"id": "a"}, {"id": "b"}]) + + res = _run(module.list_mcp()) + assert res["code"] == 0 + assert res["data"]["total"] == 2 + assert res["data"]["mcp_servers"] == [{"id": "b"}] + + monkeypatch.setattr(module, "request", SimpleNamespace(args={})) + _set_request_json(monkeypatch, module, {"mcp_ids": []}) + + def _raise_list(*_args, **_kwargs): + raise RuntimeError("list explode") + + monkeypatch.setattr(module.MCPServerService, "get_servers", _raise_list) + res = _run(module.list_mcp()) + assert res["code"] == 100 + assert "list explode" in res["message"] + + +@pytest.mark.p2 +def test_detail_not_found_success_and_exception(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + monkeypatch.setattr(module, "request", SimpleNamespace(args={"mcp_id": "mcp-1"})) + + monkeypatch.setattr(module.MCPServerService, "get_or_none", lambda **_kwargs: None) + res = module.detail() + assert res["code"] == module.RetCode.NOT_FOUND + + monkeypatch.setattr( + module.MCPServerService, + "get_or_none", + lambda **_kwargs: _DummyMCPServer(id="mcp-1", name="srv", url="http://a", server_type="sse", tenant_id="tenant_1"), + ) + res = module.detail() + assert res["code"] == 0 + assert res["data"]["id"] == "mcp-1" + + def _raise_detail(**_kwargs): + raise RuntimeError("detail explode") + + monkeypatch.setattr(module.MCPServerService, "get_or_none", _raise_detail) + res = module.detail() + assert res["code"] == 100 + assert "detail explode" in res["message"] + + +@pytest.mark.p2 +def test_create_validation_guards(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + monkeypatch.setattr(module.MCPServerService, "get_by_name_and_tenant", lambda **_kwargs: (False, None)) + + _set_request_json(monkeypatch, module, {"name": "srv", "url": "http://a", "server_type": "invalid"}) + res = _run(module.create.__wrapped__()) + assert "Unsupported MCP server type" in res["message"] + + _set_request_json(monkeypatch, module, {"name": "", "url": "http://a", "server_type": "sse"}) + res = _run(module.create.__wrapped__()) + assert "Invalid MCP name" in res["message"] + + monkeypatch.setattr(module.MCPServerService, "get_by_name_and_tenant", lambda **_kwargs: (True, object())) + _set_request_json(monkeypatch, module, {"name": "srv", "url": "http://a", "server_type": "sse"}) + res = _run(module.create.__wrapped__()) + assert "Duplicated MCP server name" in res["message"] + + monkeypatch.setattr(module.MCPServerService, "get_by_name_and_tenant", lambda **_kwargs: (False, None)) + _set_request_json(monkeypatch, module, {"name": "srv", "url": "", "server_type": "sse"}) + res = _run(module.create.__wrapped__()) + assert "Invalid url" in res["message"] + + +@pytest.mark.p2 +def test_create_service_paths(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + base_payload = { + "name": "srv", + "url": "http://server", + "server_type": "sse", + "headers": '{"Authorization": "x"}', + "variables": '{"tools": {"old": 1}, "token": "abc"}', + "timeout": "2.5", + } + + monkeypatch.setattr(module, "get_uuid", lambda: "uuid-create") + monkeypatch.setattr(module.MCPServerService, "get_by_name_and_tenant", lambda **_kwargs: (False, None)) + + _set_request_json(monkeypatch, module, dict(base_payload)) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda *_args, **_kwargs: (False, None)) + res = _run(module.create.__wrapped__()) + assert "Tenant not found" in res["message"] + + _set_request_json(monkeypatch, module, dict(base_payload)) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda *_args, **_kwargs: (True, object())) + + async def _thread_pool_tools_error(_func, _servers, _timeout): + return None, "tools error" + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_tools_error) + res = _run(module.create.__wrapped__()) + assert res["code"] == "tools error" + assert "Sorry! Data missing!" in res["message"] + + _set_request_json(monkeypatch, module, dict(base_payload)) + + async def _thread_pool_ok(_func, servers, _timeout): + return {servers[0].name: [{"name": "tool_a"}, {"invalid": True}]}, None + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_ok) + monkeypatch.setattr(module.MCPServerService, "insert", lambda **_kwargs: False) + res = _run(module.create.__wrapped__()) + assert res["code"] == "Failed to create MCP server." + assert "Sorry! Data missing!" in res["message"] + + _set_request_json(monkeypatch, module, dict(base_payload)) + monkeypatch.setattr(module.MCPServerService, "insert", lambda **_kwargs: True) + res = _run(module.create.__wrapped__()) + assert res["code"] == 0 + assert res["data"]["id"] == "uuid-create" + assert res["data"]["tenant_id"] == "tenant_1" + assert res["data"]["variables"]["tools"] == {"tool_a": {"name": "tool_a"}} + + _set_request_json(monkeypatch, module, dict(base_payload)) + + async def _thread_pool_raises(_func, _servers, _timeout): + raise RuntimeError("create explode") + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_raises) + res = _run(module.create.__wrapped__()) + assert res["code"] == 100 + assert "create explode" in res["message"] + + +@pytest.mark.p2 +def test_update_validation_guards(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + existing = _DummyMCPServer(id="mcp-1", name="srv", url="http://server", server_type="sse", tenant_id="tenant_1", variables={}, headers={}) + + _set_request_json(monkeypatch, module, {"mcp_id": "mcp-1"}) + monkeypatch.setattr(module.MCPServerService, "get_by_id", lambda _mcp_id: (False, None)) + res = _run(module.update.__wrapped__()) + assert "Cannot find MCP server" in res["message"] + + _set_request_json(monkeypatch, module, {"mcp_id": "mcp-1"}) + monkeypatch.setattr( + module.MCPServerService, + "get_by_id", + lambda _mcp_id: (True, _DummyMCPServer(id="mcp-1", name="srv", url="http://server", server_type="sse", tenant_id="other", variables={}, headers={})), + ) + res = _run(module.update.__wrapped__()) + assert "Cannot find MCP server" in res["message"] + + _set_request_json(monkeypatch, module, {"mcp_id": "mcp-1", "server_type": "invalid"}) + monkeypatch.setattr(module.MCPServerService, "get_by_id", lambda _mcp_id: (True, existing)) + res = _run(module.update.__wrapped__()) + assert "Unsupported MCP server type" in res["message"] + + _set_request_json(monkeypatch, module, {"mcp_id": "mcp-1", "name": "a" * 256}) + res = _run(module.update.__wrapped__()) + assert "Invalid MCP name" in res["message"] + + _set_request_json(monkeypatch, module, {"mcp_id": "mcp-1", "url": ""}) + res = _run(module.update.__wrapped__()) + assert "Invalid url" in res["message"] + + +@pytest.mark.p2 +def test_update_service_paths(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + existing = _DummyMCPServer( + id="mcp-1", + name="srv", + url="http://server", + server_type="sse", + tenant_id="tenant_1", + variables={"tools": {"old": {"enabled": True}}, "token": "abc"}, + headers={"Authorization": "old"}, + ) + updated = _DummyMCPServer( + id="mcp-1", + name="srv-new", + url="http://server-new", + server_type="sse", + tenant_id="tenant_1", + variables={"tools": {"tool_a": {"name": "tool_a"}}}, + headers={"Authorization": "new"}, + ) + + base_payload = { + "mcp_id": "mcp-1", + "name": "srv-new", + "url": "http://server-new", + "server_type": "sse", + "headers": '{"Authorization": "new"}', + "variables": '{"tools": {"ignore": 1}, "token": "new"}', + "timeout": "3.0", + } + + _set_request_json(monkeypatch, module, dict(base_payload)) + monkeypatch.setattr(module.MCPServerService, "get_by_id", lambda _mcp_id: (True, existing)) + + async def _thread_pool_tools_error(_func, _servers, _timeout): + return None, "update tools error" + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_tools_error) + res = _run(module.update.__wrapped__()) + assert res["code"] == "update tools error" + assert "Sorry! Data missing!" in res["message"] + + _set_request_json(monkeypatch, module, dict(base_payload)) + + async def _thread_pool_ok(_func, servers, _timeout): + return {servers[0].name: [{"name": "tool_a"}, {"bad": True}]}, None + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_ok) + monkeypatch.setattr(module.MCPServerService, "filter_update", lambda *_args, **_kwargs: False) + res = _run(module.update.__wrapped__()) + assert "Failed to updated MCP server" in res["message"] + + _set_request_json(monkeypatch, module, dict(base_payload)) + monkeypatch.setattr(module.MCPServerService, "filter_update", lambda *_args, **_kwargs: True) + + def _get_by_id_fetch_fail(_mcp_id): + if _get_by_id_fetch_fail.calls == 0: + _get_by_id_fetch_fail.calls += 1 + return True, existing + return False, None + + _get_by_id_fetch_fail.calls = 0 + monkeypatch.setattr(module.MCPServerService, "get_by_id", _get_by_id_fetch_fail) + res = _run(module.update.__wrapped__()) + assert "Failed to fetch updated MCP server" in res["message"] + + _set_request_json(monkeypatch, module, dict(base_payload)) + + def _get_by_id_success(_mcp_id): + if _get_by_id_success.calls == 0: + _get_by_id_success.calls += 1 + return True, existing + return True, updated + + _get_by_id_success.calls = 0 + monkeypatch.setattr(module.MCPServerService, "get_by_id", _get_by_id_success) + res = _run(module.update.__wrapped__()) + assert res["code"] == 0 + assert res["data"]["id"] == "mcp-1" + + _set_request_json(monkeypatch, module, dict(base_payload)) + monkeypatch.setattr(module.MCPServerService, "get_by_id", lambda _mcp_id: (True, existing)) + + async def _thread_pool_raises(_func, _servers, _timeout): + raise RuntimeError("update explode") + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_raises) + res = _run(module.update.__wrapped__()) + assert res["code"] == 100 + assert "update explode" in res["message"] + + +@pytest.mark.p2 +def test_rm_failure_success_and_exception(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + _set_request_json(monkeypatch, module, {"mcp_ids": ["a", "b"]}) + monkeypatch.setattr(module.MCPServerService, "delete_by_ids", lambda _ids: False) + res = _run(module.rm.__wrapped__()) + assert "Failed to delete MCP servers" in res["message"] + + _set_request_json(monkeypatch, module, {"mcp_ids": ["a", "b"]}) + monkeypatch.setattr(module.MCPServerService, "delete_by_ids", lambda _ids: True) + res = _run(module.rm.__wrapped__()) + assert res["code"] == 0 + assert res["data"] is True + + _set_request_json(monkeypatch, module, {"mcp_ids": ["a", "b"]}) + + def _raise_rm(_ids): + raise RuntimeError("rm explode") + + monkeypatch.setattr(module.MCPServerService, "delete_by_ids", _raise_rm) + res = _run(module.rm.__wrapped__()) + assert res["code"] == 100 + assert "rm explode" in res["message"] + + +@pytest.mark.p2 +def test_import_multiple_missing_servers_and_exception(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + _set_request_json(monkeypatch, module, {"mcpServers": {}}) + res = _run(module.import_multiple.__wrapped__()) + assert "No MCP servers provided" in res["message"] + + _set_request_json(monkeypatch, module, {"mcpServers": {"srv": {"type": "sse", "url": "http://x"}}, "timeout": "1"}) + + def _raise_import(**_kwargs): + raise RuntimeError("import explode") + + monkeypatch.setattr(module.MCPServerService, "get_by_name_and_tenant", _raise_import) + res = _run(module.import_multiple.__wrapped__()) + assert res["code"] == 100 + assert "import explode" in res["message"] + + +@pytest.mark.p2 +def test_import_multiple_mixed_results(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + payload = { + "mcpServers": { + "missing_fields": {"type": "sse"}, + "": {"type": "sse", "url": "http://empty"}, + "dup": {"type": "sse", "url": "http://dup", "authorization_token": "dup-token"}, + "tool_err": {"type": "sse", "url": "http://err"}, + "insert_fail": {"type": "sse", "url": "http://fail"}, + }, + "timeout": "3", + } + _set_request_json(monkeypatch, module, payload) + + monkeypatch.setattr(module, "get_uuid", lambda: "uuid-import") + + def _get_by_name_and_tenant(name, tenant_id): + if name == "dup" and not _get_by_name_and_tenant.first_dup_seen: + _get_by_name_and_tenant.first_dup_seen = True + return True, object() + return False, None + + _get_by_name_and_tenant.first_dup_seen = False + monkeypatch.setattr(module.MCPServerService, "get_by_name_and_tenant", _get_by_name_and_tenant) + + async def _thread_pool_exec(func, servers, _timeout): + mcp_server = servers[0] + if mcp_server.name == "tool_err": + return None, "tool call failed" + return {mcp_server.name: [{"name": "tool_a"}, {"invalid": True}]}, None + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_exec) + + def _insert(**kwargs): + return kwargs["name"] != "insert_fail" + + monkeypatch.setattr(module.MCPServerService, "insert", _insert) + + res = _run(module.import_multiple.__wrapped__()) + assert res["code"] == 0 + + results = {item["server"]: item for item in res["data"]["results"]} + assert results["missing_fields"]["success"] is False + assert "Missing required fields" in results["missing_fields"]["message"] + assert results[""]["success"] is False + assert "Invalid MCP name" in results[""]["message"] + assert results["tool_err"]["success"] is False + assert "tool call failed" in results["tool_err"]["message"] + assert results["insert_fail"]["success"] is False + assert "Failed to create MCP server" in results["insert_fail"]["message"] + assert results["dup"]["success"] is True + assert results["dup"]["new_name"] == "dup_0" + assert "Renamed from 'dup' to 'dup_0' avoid duplication" == results["dup"]["message"] + + +@pytest.mark.p2 +def test_export_multiple_missing_ids_success_and_exception(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + _set_request_json(monkeypatch, module, {"mcp_ids": []}) + res = _run(module.export_multiple.__wrapped__()) + assert "No MCP server IDs provided" in res["message"] + + _set_request_json(monkeypatch, module, {"mcp_ids": ["id1", "id2", "id3"]}) + + def _get_by_id(mcp_id): + if mcp_id == "id1": + return True, _DummyMCPServer( + id="id1", + name="srv-one", + url="http://one", + server_type="sse", + tenant_id="tenant_1", + variables={"authorization_token": "tok", "tools": {"tool_a": {"enabled": True}}}, + ) + if mcp_id == "id2": + return True, _DummyMCPServer( + id="id2", + name="srv-two", + url="http://two", + server_type="sse", + tenant_id="other", + variables={}, + ) + return False, None + + monkeypatch.setattr(module.MCPServerService, "get_by_id", _get_by_id) + res = _run(module.export_multiple.__wrapped__()) + assert res["code"] == 0 + assert list(res["data"]["mcpServers"].keys()) == ["srv-one"] + + _set_request_json(monkeypatch, module, {"mcp_ids": ["id1"]}) + + def _raise_export(_mcp_id): + raise RuntimeError("export explode") + + monkeypatch.setattr(module.MCPServerService, "get_by_id", _raise_export) + res = _run(module.export_multiple.__wrapped__()) + assert res["code"] == 100 + assert "export explode" in res["message"] + + +@pytest.mark.p2 +def test_list_tools_missing_ids_success_inner_error_outer_error_and_finally_cleanup(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + _set_request_json(monkeypatch, module, {"mcp_ids": []}) + res = _run(module.list_tools.__wrapped__()) + assert "No MCP server IDs provided" in res["message"] + + server = _DummyMCPServer( + id="id1", + name="srv-tools", + url="http://tools", + server_type="sse", + tenant_id="tenant_1", + variables={"tools": {"tool_a": {"enabled": False}}}, + ) + + _set_request_json(monkeypatch, module, {"mcp_ids": ["id1"], "timeout": "2.0"}) + monkeypatch.setattr(module.MCPServerService, "get_by_id", lambda _mcp_id: (True, server)) + + close_calls = [] + + async def _thread_pool_exec_success(func, *args): + if func is module.close_multiple_mcp_toolcall_sessions: + close_calls.append(args[0]) + return None + return func(*args) + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_exec_success) + res = _run(module.list_tools.__wrapped__()) + assert res["code"] == 0 + assert res["data"]["id1"][0]["name"] == "tool_a" + assert res["data"]["id1"][0]["enabled"] is False + assert res["data"]["id1"][1]["enabled"] is True + assert close_calls and len(close_calls[-1]) == 1 + + _set_request_json(monkeypatch, module, {"mcp_ids": ["id1"], "timeout": "2.0"}) + close_calls_inner = [] + + async def _thread_pool_exec_inner_error(func, *args): + if func is module.close_multiple_mcp_toolcall_sessions: + close_calls_inner.append(args[0]) + return None + raise RuntimeError("inner tools explode") + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_exec_inner_error) + res = _run(module.list_tools.__wrapped__()) + assert res["code"] == 102 + assert "MCP list tools error" in res["message"] + assert close_calls_inner and len(close_calls_inner[-1]) == 1 + + _set_request_json(monkeypatch, module, {"mcp_ids": ["id1"], "timeout": "2.0"}) + close_calls_outer = [] + + def _raise_get_by_id(_mcp_id): + raise RuntimeError("outer explode") + + monkeypatch.setattr(module.MCPServerService, "get_by_id", _raise_get_by_id) + + async def _thread_pool_exec_outer(func, *args): + if func is module.close_multiple_mcp_toolcall_sessions: + close_calls_outer.append(args[0]) + return None + return func(*args) + + monkeypatch.setattr(module, "thread_pool_exec", _thread_pool_exec_outer) + res = _run(module.list_tools.__wrapped__()) + assert res["code"] == 100 + assert "outer explode" in res["message"] + assert close_calls_outer + + +@pytest.mark.p2 +def test_test_tool_missing_mcp_id(monkeypatch): + module = _load_mcp_server_app(monkeypatch) + + _set_request_json(monkeypatch, module, {"mcp_id": "", "tool_name": "tool_a", "arguments": {"x": 1}}) + res = _run(module.test_tool.__wrapped__()) + assert "No MCP server ID provided" in res["message"] diff --git a/test/testcases/test_web_api/test_memory_app/test_create_memory.py b/test/testcases/test_web_api/test_memory_app/test_create_memory.py index 89e27cb8d9..d2b0f967bc 100644 --- a/test/testcases/test_web_api/test_memory_app/test_create_memory.py +++ b/test/testcases/test_web_api/test_memory_app/test_create_memory.py @@ -20,8 +20,6 @@ import pytest from test_web_api.common import create_memory from configs import INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth -from hypothesis import example, given, settings -from utils.hypothesis_utils import valid_names class TestAuthorization: @@ -42,9 +40,7 @@ class TestAuthorization: class TestMemoryCreate: @pytest.mark.p1 - @given(name=valid_names()) - @example("d" * 128) - @settings(max_examples=20) + @pytest.mark.parametrize("name", ["test_memory_name", "d" * 128]) def test_name(self, WebApiAuth, name): payload = { "name": name, @@ -79,7 +75,7 @@ class TestMemoryCreate: assert res["message"] == expected_message, res @pytest.mark.p2 - @given(name=valid_names()) + @pytest.mark.parametrize("name", ["invalid_type_name", "memory_alpha"]) def test_type_invalid(self, WebApiAuth, name): payload = { "name": name, diff --git a/test/testcases/test_web_api/test_memory_app/test_update_memory.py b/test/testcases/test_web_api/test_memory_app/test_update_memory.py index 4db2cacf5f..8b2e6c577b 100644 --- a/test/testcases/test_web_api/test_memory_app/test_update_memory.py +++ b/test/testcases/test_web_api/test_memory_app/test_update_memory.py @@ -13,14 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import re + import pytest from test_web_api.common import update_memory from configs import INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth -from hypothesis import HealthCheck, example, given, settings from utils import encode_avatar from utils.file_utils import create_image_file -from utils.hypothesis_utils import valid_names class TestAuthorization: @@ -42,15 +42,14 @@ class TestAuthorization: class TestMemoryUpdate: @pytest.mark.p1 - @given(name=valid_names()) - @example("f" * 128) - @settings(max_examples=20, suppress_health_check=[HealthCheck.function_scoped_fixture]) + @pytest.mark.parametrize("name", ["updated_memory", "f" * 128]) def test_name(self, WebApiAuth, add_memory_func, name): memory_ids = add_memory_func payload = {"name": name} res = update_memory(WebApiAuth, memory_ids[0], payload) assert res["code"] == 0, res - assert res["data"]["name"] == name, res + pattern = rf"^{re.escape(name)}(?:\(\d+\))?$" + assert re.match(pattern, res["data"]["name"]), res @pytest.mark.p2 @pytest.mark.parametrize( diff --git a/test/testcases/test_web_api/test_message_app/test_add_message.py b/test/testcases/test_web_api/test_message_app/test_add_message.py index f87b0a18c0..207c7ce115 100644 --- a/test/testcases/test_web_api/test_message_app/test_add_message.py +++ b/test/testcases/test_web_api/test_message_app/test_add_message.py @@ -71,6 +71,20 @@ Are you asking about the fruit itself, or its use in a specific context? assert message["agent_id"] == agent_id, message assert message["session_id"] == session_id, message + @pytest.mark.p2 + def test_add_message_invalid_memory_id(self, WebApiAuth): + message_payload = { + "memory_id": ["missing_memory_id"], + "agent_id": uuid.uuid4().hex, + "session_id": uuid.uuid4().hex, + "user_id": "", + "user_input": "what is pineapple?", + "agent_response": "pineapple response", + } + res = add_message(WebApiAuth, message_payload) + assert res["code"] == 500, res + assert "Some messages failed to add" in res["message"], res + @pytest.mark.usefixtures("add_empty_multiple_type_memory") class TestAddMultipleTypeMessage: diff --git a/test/testcases/test_web_api/test_message_app/test_forget_message.py b/test/testcases/test_web_api/test_message_app/test_forget_message.py index 900c321b04..fd2c5c05f7 100644 --- a/test/testcases/test_web_api/test_message_app/test_forget_message.py +++ b/test/testcases/test_web_api/test_message_app/test_forget_message.py @@ -15,8 +15,9 @@ # import random import pytest +import requests from test_web_api.common import forget_message, list_memory_message, get_message_content -from configs import INVALID_API_TOKEN +from configs import HOST_ADDRESS, INVALID_API_TOKEN, VERSION from libs.auth import RAGFlowWebApiAuth @@ -52,3 +53,17 @@ class TestForgetMessage: forgot_message_res = get_message_content(WebApiAuth, memory_id, message["message_id"]) assert forgot_message_res["code"] == 0, forgot_message_res assert forgot_message_res["data"]["forget_at"] not in ["-", ""], forgot_message_res + + @pytest.mark.p2 + def test_forget_message_invalid_memory_id(self, WebApiAuth): + res = forget_message(WebApiAuth, "missing_memory_id", 1) + assert res["code"] == 404, res + assert "not found" in res["message"].lower(), res + + @pytest.mark.p2 + def test_forget_message_invalid_message_id(self, WebApiAuth): + memory_id = self.memory_id + url = f"{HOST_ADDRESS}/api/{VERSION}/messages/{memory_id}:invalid_message_id" + res = requests.delete(url=url, headers={"Content-Type": "application/json"}, auth=WebApiAuth).json() + assert res["code"] == 500, res + assert "Internal server error" in res["message"], res diff --git a/test/testcases/test_web_api/test_message_app/test_get_message_content.py b/test/testcases/test_web_api/test_message_app/test_get_message_content.py index 35fe348d39..6ecf37b5ce 100644 --- a/test/testcases/test_web_api/test_message_app/test_get_message_content.py +++ b/test/testcases/test_web_api/test_message_app/test_get_message_content.py @@ -49,3 +49,16 @@ class TestGetMessageContent: for field in ["content", "content_embed"]: assert field in content_res["data"] assert content_res["data"][field] is not None, content_res + + @pytest.mark.p2 + def test_get_message_content_invalid_memory_id(self, WebApiAuth): + res = get_message_content(WebApiAuth, "missing_memory_id", 1) + assert res["code"] == 404, res + assert "not found" in res["message"].lower(), res + + @pytest.mark.p2 + def test_get_message_content_invalid_message_id(self, WebApiAuth): + memory_id = self.memory_id + res = get_message_content(WebApiAuth, memory_id, 999999999) + assert res["code"] == 404, res + assert "not found" in res["message"].lower(), res diff --git a/test/testcases/test_web_api/test_message_app/test_get_recent_message.py b/test/testcases/test_web_api/test_message_app/test_get_recent_message.py index 7445890f81..f7a1f846ce 100644 --- a/test/testcases/test_web_api/test_message_app/test_get_recent_message.py +++ b/test/testcases/test_web_api/test_message_app/test_get_recent_message.py @@ -66,3 +66,15 @@ class TestGetRecentMessage: for message in res["data"]: assert message["session_id"] == session_id, message + @pytest.mark.p2 + def test_get_recent_messages_missing_memory_id(self, WebApiAuth): + res = get_recent_message(WebApiAuth, params={}) + assert res["code"] == 101, res + assert "memory_ids is required" in res["message"], res + + @pytest.mark.p2 + def test_get_recent_messages_csv_memory_ids(self, WebApiAuth): + memory_id = self.memory_id + res = get_recent_message(WebApiAuth, params={"memory_id": f"{memory_id},{memory_id}"}) + assert res["code"] == 0, res + assert isinstance(res["data"], list), res diff --git a/test/testcases/test_web_api/test_message_app/test_message_routes_unit.py b/test/testcases/test_web_api/test_message_app/test_message_routes_unit.py new file mode 100644 index 0000000000..f4641ed469 --- /dev/null +++ b/test/testcases/test_web_api/test_message_app/test_message_routes_unit.py @@ -0,0 +1,151 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import asyncio +import importlib.util +import inspect +import sys +from copy import deepcopy +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + + return decorator + + +class _AwaitableValue: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _co(): + return self._value + + return _co().__await__() + + +class _DummyArgs(dict): + def getlist(self, key): + value = self.get(key) + if value is None: + return [] + if isinstance(value, list): + return value + return [value] + + +class _DummyMemoryApiService: + async def add_message(self, *_args, **_kwargs): + return True, "ok" + + async def get_messages(self, *_args, **_kwargs): + return [] + + +def _run(coro): + return asyncio.run(coro) + + +def _load_memory_routes_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + apps_mod = ModuleType("api.apps") + apps_mod.__path__ = [str(repo_root / "api" / "apps")] + apps_mod.current_user = SimpleNamespace(id="user-1") + apps_mod.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + + services_mod = ModuleType("api.apps.services") + services_mod.memory_api_service = _DummyMemoryApiService() + monkeypatch.setitem(sys.modules, "api.apps.services", services_mod) + + module_name = "test_message_routes_unit_module" + module_path = repo_root / "api" / "apps" / "restful_apis" / "memory_api.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _set_request_json(monkeypatch, module, payload): + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(deepcopy(payload))) + + +@pytest.mark.p2 +def test_add_message_partial_failure_branch(monkeypatch): + module = _load_memory_routes_module(monkeypatch) + + _set_request_json( + monkeypatch, + module, + { + "memory_id": ["memory-1"], + "agent_id": "agent-1", + "session_id": "session-1", + "user_input": "hello", + "agent_response": "world", + }, + ) + + async def _add_message(_memory_ids, _message_dict): + return False, "cannot enqueue" + + monkeypatch.setattr(module.memory_api_service, "add_message", _add_message) + + res = _run(inspect.unwrap(module.add_message)()) + assert res["code"] == module.RetCode.SERVER_ERROR, res + assert "Some messages failed to add" in res["message"], res + + +@pytest.mark.p2 +def test_get_messages_csv_and_missing_memory_ids(monkeypatch): + module = _load_memory_routes_module(monkeypatch) + + monkeypatch.setattr(module, "request", SimpleNamespace(args=_DummyArgs({}))) + res = _run(inspect.unwrap(module.get_messages)()) + assert res["code"] == module.RetCode.ARGUMENT_ERROR, res + assert "memory_ids is required." in res["message"], res + + monkeypatch.setattr( + module, + "request", + SimpleNamespace(args=_DummyArgs({"memory_id": "m1,m2", "agent_id": "a1", "session_id": "s1", "limit": "5"})), + ) + + async def _get_messages(memory_ids, agent_id, session_id, limit): + assert memory_ids == ["m1", "m2"] + assert agent_id == "a1" + assert session_id == "s1" + assert limit == 5 + return [{"message_id": 1}] + + monkeypatch.setattr(module.memory_api_service, "get_messages", _get_messages) + res = _run(inspect.unwrap(module.get_messages)()) + assert res["code"] == module.RetCode.SUCCESS, res + assert isinstance(res["data"], list), res diff --git a/test/testcases/test_web_api/test_message_app/test_search_message.py b/test/testcases/test_web_api/test_message_app/test_search_message.py index 0c82bc5bef..eafdde6913 100644 --- a/test/testcases/test_web_api/test_message_app/test_search_message.py +++ b/test/testcases/test_web_api/test_message_app/test_search_message.py @@ -80,3 +80,23 @@ class TestSearchMessage: assert res["code"] == 0, res assert len(res["data"]) > 0 assert len(res["data"]) <= params["top_n"] + + @pytest.mark.p2 + def test_query_missing_query(self, WebApiAuth): + memory_id = self.memory_id + res = search_message(WebApiAuth, {"memory_id": memory_id}) + assert res["code"] in [100, 500], res + + @pytest.mark.p2 + def test_query_missing_memory_id(self, WebApiAuth): + res = search_message(WebApiAuth, {"query": "what is coriander"}) + assert res["code"] == 0, res + assert isinstance(res["data"], list), res + + @pytest.mark.p2 + def test_query_with_csv_memory_ids(self, WebApiAuth): + memory_id = self.memory_id + query = "Coriander is a versatile herb." + res = search_message(WebApiAuth, {"memory_id": f"{memory_id},{memory_id}", "query": query}) + assert res["code"] == 0, res + assert isinstance(res["data"], list), res diff --git a/test/testcases/test_web_api/test_message_app/test_update_message_status.py b/test/testcases/test_web_api/test_message_app/test_update_message_status.py index 50e9df3ad8..d1df6a159d 100644 --- a/test/testcases/test_web_api/test_message_app/test_update_message_status.py +++ b/test/testcases/test_web_api/test_message_app/test_update_message_status.py @@ -16,9 +16,11 @@ import random import pytest +import requests from test_web_api.common import update_message_status, list_memory_message, get_message_content from configs import INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth +from configs import HOST_ADDRESS, VERSION class TestAuthorization: @@ -73,3 +75,34 @@ class TestUpdateMessageStatus: res = get_message_content(WebApiAuth, memory_id, message["message_id"]) assert res["code"] == 0, res assert res["data"]["status"], res + + @pytest.mark.p2 + def test_update_invalid_status_type(self, WebApiAuth): + memory_id = self.memory_id + list_res = list_memory_message(WebApiAuth, memory_id) + assert list_res["code"] == 0, list_res + message_id = list_res["data"]["messages"]["message_list"][0]["message_id"] + + url = f"{HOST_ADDRESS}/api/{VERSION}/messages/{memory_id}:{message_id}" + res = requests.put(url=url, headers={"Content-Type": "application/json"}, auth=WebApiAuth, json={"status": "false"}).json() + assert res["code"] == 101, res + assert "Status must be a boolean." in res["message"], res + + @pytest.mark.p2 + def test_update_invalid_memory_id(self, WebApiAuth): + res = update_message_status(WebApiAuth, "missing_memory_id", 1, False) + assert res["code"] == 404, res + assert "not found" in res["message"].lower(), res + + @pytest.mark.p2 + def test_update_invalid_message_id(self, WebApiAuth): + memory_id = self.memory_id + url = f"{HOST_ADDRESS}/api/{VERSION}/messages/{memory_id}:invalid_message_id" + res = requests.put( + url=url, + headers={"Content-Type": "application/json"}, + auth=WebApiAuth, + json={"status": True}, + ).json() + assert res["code"] == 500, res + assert "Internal server error" in res["message"], res diff --git a/test/testcases/test_web_api/test_plugin_app/test_llm_tools.py b/test/testcases/test_web_api/test_plugin_app/test_llm_tools.py index 9b5cec085c..6f1316664d 100644 --- a/test/testcases/test_web_api/test_plugin_app/test_llm_tools.py +++ b/test/testcases/test_web_api/test_plugin_app/test_llm_tools.py @@ -13,6 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + import pytest from common import plugin_llm_tools from configs import INVALID_API_TOKEN @@ -40,3 +45,54 @@ class TestPluginTools: res = plugin_llm_tools(WebApiAuth) assert res["code"] == 0, res assert isinstance(res["data"], list), res + + +class _DummyManager: + def route(self, *_args, **_kwargs): + def decorator(func): + return func + return decorator + + +def _load_plugin_app(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + stub_apps = ModuleType("api.apps") + stub_apps.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", stub_apps) + + stub_plugin = ModuleType("agent.plugin") + + class _StubGlobalPluginManager: + @staticmethod + def get_llm_tools(): + return [] + + stub_plugin.GlobalPluginManager = _StubGlobalPluginManager + monkeypatch.setitem(sys.modules, "agent.plugin", stub_plugin) + + module_path = Path(__file__).resolve().parents[4] / "api" / "apps" / "plugin_app.py" + spec = importlib.util.spec_from_file_location("test_plugin_app_unit", module_path) + module = importlib.util.module_from_spec(spec) + module.manager = _DummyManager() + spec.loader.exec_module(module) + return module + + +@pytest.mark.p2 +def test_llm_tools_metadata_shape_unit(monkeypatch): + module = _load_plugin_app(monkeypatch) + + class _DummyTool: + def get_metadata(self): + return {"name": "dummy", "description": "test"} + + monkeypatch.setattr(module.GlobalPluginManager, "get_llm_tools", staticmethod(lambda: [_DummyTool()])) + res = module.llm_tools() + assert res["code"] == 0 + assert isinstance(res["data"], list) + assert res["data"][0]["name"] == "dummy" + assert res["data"][0]["description"] == "test" diff --git a/test/testcases/test_web_api/test_system_app/test_apps_init_unit.py b/test/testcases/test_web_api/test_system_app/test_apps_init_unit.py new file mode 100644 index 0000000000..cfd79ce879 --- /dev/null +++ b/test/testcases/test_web_api/test_system_app/test_apps_init_unit.py @@ -0,0 +1,242 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import asyncio +import importlib.util +import logging +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +from werkzeug.exceptions import Unauthorized as WerkzeugUnauthorized + + +class _DummyAPIToken: + @staticmethod + def query(**_kwargs): + return [] + + +class _DummyUserService: + @staticmethod + def query(**_kwargs): + return [] + + +def _run(coro): + return asyncio.run(coro) + + +def _load_apps_module(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + common_pkg = ModuleType("common") + common_pkg.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common_pkg) + + settings_mod = ModuleType("common.settings") + settings_mod.SECRET_KEY = "test-secret-key" + settings_mod.init_settings = lambda: None + settings_mod.decrypt_database_config = lambda name=None: {} + monkeypatch.setitem(sys.modules, "common.settings", settings_mod) + common_pkg.settings = settings_mod + + db_models_mod = ModuleType("api.db.db_models") + db_models_mod.APIToken = _DummyAPIToken + db_models_mod.close_connection = lambda: None + monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod) + + services_mod = ModuleType("api.db.services") + services_mod.UserService = _DummyUserService + monkeypatch.setitem(sys.modules, "api.db.services", services_mod) + + commands_mod = ModuleType("api.utils.commands") + commands_mod.register_commands = lambda _app: None + monkeypatch.setitem(sys.modules, "api.utils.commands", commands_mod) + + api_utils_mod = ModuleType("api.utils.api_utils") + + def _get_json_result(code=0, message="success", data=None): + return {"code": code, "message": message, "data": data} + + def _server_error_response(error): + return {"code": 100, "message": repr(error)} + + api_utils_mod.get_json_result = _get_json_result + api_utils_mod.server_error_response = _server_error_response + monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod) + + module_name = "test_apps_init_unit_module" + module_path = repo_root / "api" / "apps" / "__init__.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + + monkeypatch.setattr(Path, "glob", lambda self, _pattern: []) + spec.loader.exec_module(module) + return module.app, module + + +@pytest.mark.p2 +def test_module_init_and_unauthorized_message_variants(monkeypatch): + _quart_app, apps_module = _load_apps_module(monkeypatch) + + assert apps_module.client_urls_prefix == [] + + class _BrokenRepr: + def __repr__(self): + raise RuntimeError("repr explode") + + class _ExactUnauthorizedRepr: + def __repr__(self): + return apps_module.UNAUTHORIZED_MESSAGE + + class _Unauthorized401Repr: + def __repr__(self): + return "Unauthorized 401 from upstream" + + class _OtherRepr: + def __repr__(self): + return "Forbidden 403" + + assert apps_module._unauthorized_message(None) == apps_module.UNAUTHORIZED_MESSAGE + assert apps_module._unauthorized_message(_BrokenRepr()) == apps_module.UNAUTHORIZED_MESSAGE + assert apps_module._unauthorized_message(_ExactUnauthorizedRepr()) == apps_module.UNAUTHORIZED_MESSAGE + assert apps_module._unauthorized_message(_Unauthorized401Repr()) == "Unauthorized 401 from upstream" + assert apps_module._unauthorized_message(_OtherRepr()) == apps_module.UNAUTHORIZED_MESSAGE + + +@pytest.mark.p2 +def test_load_user_token_edge_cases(monkeypatch): + quart_app, apps_module = _load_apps_module(monkeypatch) + + user_with_empty_token = SimpleNamespace(email="empty@example.com", access_token="") + + async def _case(): + async with quart_app.test_request_context("/", headers={"Authorization": "token"}): + monkeypatch.setattr(apps_module.Serializer, "loads", lambda _self, _auth: "") + assert apps_module._load_user() is None + + async with quart_app.test_request_context("/", headers={"Authorization": "token"}): + monkeypatch.setattr(apps_module.Serializer, "loads", lambda _self, _auth: "short-token") + assert apps_module._load_user() is None + + async with quart_app.test_request_context("/", headers={"Authorization": "token"}): + monkeypatch.setattr(apps_module.Serializer, "loads", lambda _self, _auth: "a" * 32) + monkeypatch.setattr(apps_module.UserService, "query", lambda **_kwargs: [user_with_empty_token]) + assert apps_module._load_user() is None + + _run(_case()) + + +@pytest.mark.p2 +def test_load_user_api_token_fallback_and_fallback_exception(monkeypatch, caplog): + quart_app, apps_module = _load_apps_module(monkeypatch) + + def _raise_decode(_self, _auth): + raise RuntimeError("decode failed") + + monkeypatch.setattr(apps_module.Serializer, "loads", _raise_decode) + + fallback_user_empty_token = SimpleNamespace(email="fallback@example.com", access_token="") + + async def _case(): + monkeypatch.setattr(apps_module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) + monkeypatch.setattr(apps_module.UserService, "query", lambda **_kwargs: [fallback_user_empty_token]) + async with quart_app.test_request_context("/", headers={"Authorization": "Bearer api-token"}): + assert apps_module._load_user() is None + + def _raise_api_token(**_kwargs): + raise RuntimeError("api token fallback failed") + + monkeypatch.setattr(apps_module.APIToken, "query", _raise_api_token) + async with quart_app.test_request_context("/", headers={"Authorization": "Bearer api-token"}): + with caplog.at_level(logging.WARNING): + assert apps_module._load_user() is None + + _run(_case()) + assert "api token fallback failed" in caplog.text + + +@pytest.mark.p2 +def test_login_required_timing_and_login_user_inactive(monkeypatch, caplog): + quart_app, apps_module = _load_apps_module(monkeypatch) + + monkeypatch.setenv("RAGFLOW_API_TIMING", "1") + monkeypatch.setattr(apps_module, "current_user", SimpleNamespace(id="tenant-1")) + + @apps_module.login_required + async def _timed_handler(): + return {"ok": True} + + async def _case(): + async with quart_app.test_request_context("/timed"): + with caplog.at_level(logging.INFO): + assert await _timed_handler() == {"ok": True} + + inactive_user = SimpleNamespace(id="user-1", is_active=False) + assert apps_module.login_user(inactive_user) is False + + _run(_case()) + assert "api_timing login_required" in caplog.text + + +@pytest.mark.p2 +def test_logout_user_not_found_and_unauthorized_handlers(monkeypatch): + quart_app, apps_module = _load_apps_module(monkeypatch) + + async def _case(): + async with quart_app.test_request_context("/logout", headers={"Cookie": "remember_token=abc"}): + from quart import session + + session["_user_id"] = "user-1" + session["_fresh"] = True + session["_id"] = "session-id" + session["_remember_seconds"] = 5 + + assert apps_module.logout_user() is True + assert "_user_id" not in session + assert "_fresh" not in session + assert "_id" not in session + assert session.get("_remember") == "clear" + assert "_remember_seconds" not in session + + async with quart_app.test_request_context("/missing/path"): + not_found_resp, status = await apps_module.not_found(RuntimeError("missing")) + assert status == apps_module.RetCode.NOT_FOUND + payload = await not_found_resp.get_json() + assert payload["code"] == apps_module.RetCode.NOT_FOUND + assert payload["error"] == "Not Found" + assert "Not Found:" in payload["message"] + + async with quart_app.test_request_context("/protected"): + @apps_module.login_required + async def _protected(): + return {"ok": True} + + monkeypatch.setattr(apps_module, "current_user", None) + with pytest.raises(apps_module.QuartAuthUnauthorized) as exc_info: + await _protected() + + quart_payload, quart_status = await apps_module.unauthorized_quart_auth(exc_info.value) + assert quart_status == apps_module.RetCode.UNAUTHORIZED + assert quart_payload["code"] == apps_module.RetCode.UNAUTHORIZED + + werk_payload, werk_status = await apps_module.unauthorized_werkzeug(WerkzeugUnauthorized("Unauthorized 401")) + assert werk_status == apps_module.RetCode.UNAUTHORIZED + assert werk_payload["code"] == apps_module.RetCode.UNAUTHORIZED + + _run(_case())