mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-29 04:08:12 +08:00
tests: improve RAGFlow coverage based on Codecov report (#13200)
### What problem does this PR solve? Codecov’s coverage report shows that several RAGFlow code paths are currently untested or under-tested. This makes it easier for regressions to slip in during refactors and feature work. This PR adds targeted automated tests to cover the files and branches highlighted by Codecov, improving confidence in core behavior while keeping runtime functionality unchanged. ### Type of change - [x] Other (please describe): Test coverage improvement (adds/extends unit and integration tests to address Codecov-reported gaps)
This commit is contained in:
6ba3i
@@ -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"
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
1055
test/testcases/test_http_api/test_file_app/test_file_routes.py
Normal file
1055
test/testcases/test_http_api/test_file_app/test_file_routes.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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"]
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user