feat(api): add unified index API and dataset management endpoints (#14222)

### What problem does this PR solve?

## Summary

Refactor the dataset API layer into a clean service/REST separation
pattern, add a unified `/index` API for graph/raptor/mindmap operations,
and introduce several new dataset management endpoints with full test
coverage.

## Changes

### Service Layer (`dataset_api_service.py`)

- Added `trace_index(dataset_id, tenant_id, index_type)` — unified trace
function for all index types
- Added `run_index`, `delete_index` service functions
- Added `get_dataset`, `get_ingestion_summary`, `list_ingestion_logs`,
`get_ingestion_log`
- Added `run_embedding`, `list_tags`, `aggregate_tags`, `delete_tags`,
`rename_tag`
- Added `get_flattened_metadata`, `get_auto_metadata`,
`update_auto_metadata`

### REST API Layer (`dataset_api.py`)

**New unified routes:**

| Method | Route | Description |
|--------|-------|-------------|
| POST | `/datasets/<id>/index?type=graph\|raptor\|mindmap` | Run index
task |
| GET | `/datasets/<id>/index?type=graph\|raptor\|mindmap` | Trace index
task |
| DELETE | `/datasets/<id>/<index_type>` | Delete index |
| GET | `/datasets/<id>` | Get dataset details |
| GET | `/datasets/<id>/ingestions/summary` | Ingestion summary |
| GET | `/datasets/<id>/ingestions` | List ingestion logs |
| GET | `/datasets/<id>/ingestions/<log_id>` | Get single ingestion log
|
| POST | `/datasets/<id>/embedding` | Run embedding |
| GET | `/datasets/<id>/tags` | List tags |
| GET | `/datasets/tags/aggregation` | Aggregate tags across datasets |
| DELETE | `/datasets/<id>/tags` | Delete tags |
| PUT | `/datasets/<id>/tags` | Rename tag |
| GET | `/datasets/metadata/flattened` | Get flattened metadata |
| GET/PUT | `/datasets/<id>/metadata/config` | New metadata config path
|

**Removed routes (replaced by unified `/index`):**

- `POST /datasets/<id>/mindmap`
- `GET /datasets/<id>/mindmap`

**Preserved legacy routes (backward compatibility):**

- `/run_graphrag`, `/trace_graphrag`, `/run_raptor`, `/trace_raptor`
- `/auto_metadata` GET/PUT

### Test Suite

- Updated `common.py` helpers: added `trace_index`, removed
`run_mindmap`/`trace_mindmap`
- Added 7 new test files with 39 test cases total:

| Test File | Cases |
|-----------|-------|
| `test_get_dataset.py` | 4 |
| `test_ingestion_summary.py` | 2 |
| `test_ingestion_logs.py` | 5 |
| `test_index_api.py` | 14 |
| `test_embedding.py` | 2 |
| `test_tags.py` | 8 |
| `test_flattened_metadata.py` | 4 |

- Deleted `test_mindmap_tasks.py` (covered by unified index tests)

## Design Decisions

1. **Unified `/index?type=...`** — single endpoint replaces 3 separate
route pairs for graph/raptor/mindmap
2. **Backward compatibility** — old routes (`/run_graphrag`,
`/run_raptor`, `/auto_metadata`) preserved alongside new paths
3. **`_VALID_INDEX_TYPES = {"graph", "raptor", "mindmap"}`** — input
validation via constant set
4. **`_INDEX_TYPE_TO_TASK_ID_FIELD`** — maps index type to KB model task
ID field for clean dispatch

## Files Changed

- `api/apps/restful_apis/dataset_api.py`
- `api/apps/services/dataset_api_service.py`
- `sdk/python/ragflow_sdk/modules/dataset.py`
- `test/testcases/test_http_api/common.py`
- `test/testcases/test_http_api/test_dataset_management/` (7 new files)
### Type of change

- [x] New Feature (non-breaking change which adds functionality)
- [x] Refactoring

---------

Signed-off-by: noob <yixiao121314@outlook.com>
This commit is contained in:
euvre
2026-04-27 01:38:01 +00:00
committed by GitHub
parent fb95136f39
commit 4dcc42e0e1
51 changed files with 1765 additions and 4381 deletions

View File

@@ -194,14 +194,14 @@ class TestChunksRetrieval:
100,
4,
"must be greater than 0",
marks=pytest.mark.skipif(os.getenv("DOC_ENGINE") in ["infinity", "opensearch"], reason="Infinity"),
marks=pytest.mark.skip(reason="Web API does not validate top_k"),
),
pytest.param(
{"top_k": -1},
100,
4,
"3014",
marks=pytest.mark.skipif(os.getenv("DOC_ENGINE") in [None, "opensearch", "elasticsearch"], reason="elasticsearch"),
marks=pytest.mark.skip(reason="Web API does not validate top_k"),
),
pytest.param(
{"top_k": "a"},

View File

@@ -25,7 +25,6 @@ from utils.file_utils import create_txt_file
HEADERS = {"Content-Type": "application/json"}
KB_APP_URL = f"/{VERSION}/kb"
DATASETS_URL = f"/api/{VERSION}/datasets"
DOCUMENT_APP_URL = f"/{VERSION}/document"
CHUNK_APP_URL = f"/{VERSION}/chunk"
@@ -207,49 +206,41 @@ def delete_datasets(auth, payload=None, *, headers=HEADERS, data=None):
return res.json()
def detail_kb(auth, params=None, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{KB_APP_URL}/detail", headers=headers, auth=auth, params=params)
def detail_kb(auth, dataset_id, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}", headers=headers, auth=auth)
return res.json()
def kb_get_meta(auth, params=None, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{KB_APP_URL}/get_meta", headers=headers, auth=auth, params=params)
def kb_get_meta(auth, dataset_ids, *, headers=HEADERS):
params = {"dataset_ids": dataset_ids}
res = requests.get(url=f"{HOST_ADDRESS}{DATASETS_URL}/metadata/flattened", headers=headers, auth=auth, params=params)
return res.json()
def kb_basic_info(auth, params=None, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{KB_APP_URL}/basic_info", headers=headers, auth=auth, params=params)
def kb_basic_info(auth, dataset_id, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/ingestions/summary", headers=headers, auth=auth)
return res.json()
def kb_update_metadata_setting(auth, payload=None, *, headers=HEADERS, data=None):
res = requests.post(url=f"{HOST_ADDRESS}{KB_APP_URL}/update_metadata_setting", headers=headers, auth=auth, json=payload, data=data)
def kb_update_metadata_setting(auth, dataset_id, payload=None, *, headers=HEADERS, data=None):
res = requests.put(url=f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/metadata/config", headers=headers, auth=auth, json=payload, data=data)
return res.json()
def kb_list_pipeline_logs(auth, params=None, payload=None, *, headers=HEADERS, data=None):
if payload is None:
payload = {}
res = requests.post(url=f"{HOST_ADDRESS}{KB_APP_URL}/list_pipeline_logs", headers=headers, auth=auth, params=params, json=payload, data=data)
def kb_list_pipeline_logs(auth, dataset_id, params=None, *, headers=HEADERS):
url = f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/ingestions"
res = requests.get(url=url, headers=headers, auth=auth, params=params)
return res.json()
def kb_list_pipeline_dataset_logs(auth, params=None, payload=None, *, headers=HEADERS, data=None):
if payload is None:
payload = {}
res = requests.post(url=f"{HOST_ADDRESS}{KB_APP_URL}/list_pipeline_dataset_logs", headers=headers, auth=auth, params=params, json=payload, data=data)
def kb_list_pipeline_dataset_logs(auth, dataset_id, params=None, *, headers=HEADERS):
url = f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/ingestions"
res = requests.get(url=url, headers=headers, auth=auth, params=params)
return res.json()
def kb_delete_pipeline_logs(auth, params=None, payload=None, *, headers=HEADERS, data=None):
if payload is None:
payload = {}
res = requests.post(url=f"{HOST_ADDRESS}{KB_APP_URL}/delete_pipeline_logs", headers=headers, auth=auth, params=params, json=payload, data=data)
return res.json()
def kb_pipeline_log_detail(auth, params=None, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{KB_APP_URL}/pipeline_log_detail", headers=headers, auth=auth, params=params)
def kb_pipeline_log_detail(auth, dataset_id, log_id, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/ingestions/{log_id}", headers=headers, auth=auth)
return res.json()
@@ -269,57 +260,24 @@ def delete_knowledge_graph(auth, dataset_id, payload=None):
return res.json()
def run_graphrag(auth, dataset_id, payload=None):
url = f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/run_graphrag"
res = requests.post(url=url, headers=HEADERS, auth=auth, json=payload)
def list_tags_from_kbs(auth, dataset_ids, *, headers=HEADERS):
params = {"dataset_ids": dataset_ids}
res = requests.get(url=f"{HOST_ADDRESS}{DATASETS_URL}/tags/aggregation", headers=headers, auth=auth, params=params)
return res.json()
def trace_graphrag(auth, dataset_id, params=None):
url = f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/trace_graphrag"
res = requests.get(url=url, headers=HEADERS, auth=auth, params=params)
return res.json()
def run_raptor(auth, dataset_id, payload=None):
url = f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/run_raptor"
res = requests.post(url=url, headers=HEADERS, auth=auth, json=payload)
return res.json()
def trace_raptor(auth, dataset_id, params=None):
url = f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/trace_raptor"
res = requests.get(url=url, headers=HEADERS, auth=auth, params=params)
return res.json()
def kb_run_mindmap(auth, payload=None, *, headers=HEADERS, data=None):
res = requests.post(url=f"{HOST_ADDRESS}{KB_APP_URL}/run_mindmap", headers=headers, auth=auth, json=payload, data=data)
return res.json()
def kb_trace_mindmap(auth, params=None, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{KB_APP_URL}/trace_mindmap", headers=headers, auth=auth, params=params)
return res.json()
def list_tags_from_kbs(auth, params=None, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{KB_APP_URL}/tags", headers=headers, auth=auth, params=params)
return res.json()
def list_tags(auth, dataset_id, params=None, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{KB_APP_URL}/{dataset_id}/tags", headers=headers, auth=auth, params=params)
def list_tags(auth, dataset_id, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/tags", headers=headers, auth=auth)
return res.json()
def rm_tags(auth, dataset_id, payload=None, *, headers=HEADERS, data=None):
res = requests.post(url=f"{HOST_ADDRESS}{KB_APP_URL}/{dataset_id}/rm_tags", headers=headers, auth=auth, json=payload, data=data)
res = requests.delete(url=f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/tags", headers=headers, auth=auth, json=payload, data=data)
return res.json()
def rename_tags(auth, dataset_id, payload=None, *, headers=HEADERS, data=None):
res = requests.post(url=f"{HOST_ADDRESS}{KB_APP_URL}/{dataset_id}/rename_tag", headers=headers, auth=auth, json=payload, data=data)
res = requests.put(url=f"{HOST_ADDRESS}{DATASETS_URL}/{dataset_id}/tags", headers=headers, auth=auth, json=payload, data=data)
return res.json()

View File

@@ -142,6 +142,12 @@ def _load_dataset_module(monkeypatch):
api_pkg.__path__ = [str(repo_root / "api")]
monkeypatch.setitem(sys.modules, "api", api_pkg)
api_constants_mod = ModuleType("api.constants")
api_constants_mod.DATASET_NAME_LIMIT = 128
api_constants_mod.FILE_NAME_LEN_LIMIT = 255
monkeypatch.setitem(sys.modules, "api.constants", api_constants_mod)
api_pkg.constants = api_constants_mod
utils_pkg = ModuleType("api.utils")
utils_pkg.__path__ = [str(repo_root / "api" / "utils")]
monkeypatch.setitem(sys.modules, "api.utils", utils_pkg)
@@ -161,6 +167,7 @@ def _load_dataset_module(monkeypatch):
db_pkg = ModuleType("api.db")
db_pkg.__path__ = []
db_pkg.FileType = SimpleNamespace()
monkeypatch.setitem(sys.modules, "api.db", db_pkg)
api_pkg.db = db_pkg
@@ -313,8 +320,14 @@ def _load_dataset_module(monkeypatch):
def get_by_ids(_ids):
return []
class _StubUserTenantService:
@staticmethod
def get_tenants_by_user_id(_user_id):
return []
user_service_mod.TenantService = _StubTenantService
user_service_mod.UserService = _StubUserService
user_service_mod.UserTenantService = _StubUserTenantService
monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod)
services_pkg.user_service = user_service_mod
@@ -662,143 +675,115 @@ def test_list_knowledge_graph_delete_kg_matrix_unit(monkeypatch):
@pytest.mark.p3
def test_run_trace_graphrag_matrix_unit(monkeypatch):
def test_run_index_matrix_unit(monkeypatch):
module = _load_dataset_module(monkeypatch)
warnings = []
monkeypatch.setattr(module.logging, "warning", lambda msg, *_args, **_kwargs: warnings.append(msg))
res = _run(inspect.unwrap(module.run_graphrag)("tenant-1", ""))
assert 'Dataset ID' in res["message"], res
# Invalid index type
_set_request_args(monkeypatch, module, {"type": "invalid"})
res = _run(inspect.unwrap(module.run_index)("tenant-1", "kb-1"))
assert "Invalid index type" in res["message"], res
# Missing dataset ID
_set_request_args(monkeypatch, module, {"type": "graph"})
res = _run(inspect.unwrap(module.run_index)("tenant-1", ""))
assert "Dataset ID" in res["message"], res
# No authorization
_set_request_args(monkeypatch, module, {"type": "graph"})
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False)
res = _run(inspect.unwrap(module.run_graphrag)("tenant-1", "kb-1"))
res = _run(inspect.unwrap(module.run_index)("tenant-1", "kb-1"))
assert res["code"] == module.RetCode.DATA_ERROR, res
# Invalid dataset ID
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True)
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None))
res = _run(inspect.unwrap(module.run_graphrag)("tenant-1", "kb-1"))
res = _run(inspect.unwrap(module.run_index)("tenant-1", "kb-1"))
assert "Invalid Dataset ID" in res["message"], res
# Stale graphrag task + successful re-queue
stale_kb = _KB(kb_id="kb-1", graphrag_task_id="task-old")
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, stale_kb))
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (False, None))
monkeypatch.setattr(module.DocumentService, "get_by_kb_id", lambda **_kwargs: ([{"id": "doc-1"}], 1))
monkeypatch.setattr(module.dataset_api_service, "queue_raptor_o_graphrag_tasks", lambda **_kwargs: "task-new")
monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: True)
res = _run(inspect.unwrap(module.run_graphrag)("tenant-1", "kb-1"))
_set_request_args(monkeypatch, module, {"type": "graph"})
res = _run(inspect.unwrap(module.run_index)("tenant-1", "kb-1"))
assert res["code"] == module.RetCode.SUCCESS, res
assert any("GraphRAG" in msg for msg in warnings), warnings
assert any("Graph" in msg for msg in warnings), warnings
# Task already running
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (True, SimpleNamespace(progress=0)))
res = _run(inspect.unwrap(module.run_graphrag)("tenant-1", "kb-1"))
res = _run(inspect.unwrap(module.run_index)("tenant-1", "kb-1"))
assert "already running" in res["message"], res
# Successful raptor run with save warning
warnings.clear()
queue_calls = {}
no_task_kb = _KB(kb_id="kb-1", graphrag_task_id="")
no_task_kb = _KB(kb_id="kb-1", raptor_task_id="")
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, no_task_kb))
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (False, None))
monkeypatch.setattr(module.DocumentService, "get_by_kb_id", lambda **_kwargs: ([{"id": "doc-1"}, {"id": "doc-2"}], 2))
queue_calls = {}
def _queue(**kwargs):
queue_calls.update(kwargs)
return "queued-id"
return "queued-raptor"
monkeypatch.setattr(module.dataset_api_service, "queue_raptor_o_graphrag_tasks", _queue)
monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: False)
res = _run(inspect.unwrap(module.run_graphrag)("tenant-1", "kb-1"))
_set_request_args(monkeypatch, module, {"type": "raptor"})
res = _run(inspect.unwrap(module.run_index)("tenant-1", "kb-1"))
assert res["code"] == module.RetCode.SUCCESS, res
assert res["data"]["graphrag_task_id"] == "queued-id", res
assert res["data"]["task_id"] == "queued-raptor", res
assert queue_calls["doc_ids"] == ["doc-1", "doc-2"], queue_calls
assert any("Cannot save graphrag_task_id" in msg for msg in warnings), warnings
res = inspect.unwrap(module.trace_graphrag)("tenant-1", "")
assert 'Dataset ID' in res["message"], res
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False)
res = inspect.unwrap(module.trace_graphrag)("tenant-1", "kb-1")
assert res["code"] == module.RetCode.DATA_ERROR, res
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True)
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None))
res = inspect.unwrap(module.trace_graphrag)("tenant-1", "kb-1")
assert "Invalid Dataset ID" in res["message"], res
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _KB(kb_id="kb-1", graphrag_task_id="task-1")))
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (False, None))
res = inspect.unwrap(module.trace_graphrag)("tenant-1", "kb-1")
assert res["code"] == module.RetCode.SUCCESS, res
assert res["data"] == {}, res
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (True, SimpleNamespace(to_dict=lambda: {"id": _task_id, "progress": 1})))
res = inspect.unwrap(module.trace_graphrag)("tenant-1", "kb-1")
assert res["code"] == module.RetCode.SUCCESS, res
assert res["data"]["id"] == "task-1", res
assert any("Cannot save" in msg for msg in warnings), warnings
@pytest.mark.p3
def test_run_trace_raptor_matrix_unit(monkeypatch):
def test_trace_index_matrix_unit(monkeypatch):
module = _load_dataset_module(monkeypatch)
warnings = []
monkeypatch.setattr(module.logging, "warning", lambda msg, *_args, **_kwargs: warnings.append(msg))
# Invalid index type
_set_request_args(monkeypatch, module, {"type": "invalid"})
res = inspect.unwrap(module.trace_index)("tenant-1", "kb-1")
assert "Invalid index type" in res["message"], res
res = _run(inspect.unwrap(module.run_raptor)("tenant-1", ""))
assert 'Dataset ID' in res["message"], res
# Missing dataset ID
_set_request_args(monkeypatch, module, {"type": "graph"})
res = inspect.unwrap(module.trace_index)("tenant-1", "")
assert "Dataset ID" in res["message"], res
# No authorization
_set_request_args(monkeypatch, module, {"type": "graph"})
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False)
res = _run(inspect.unwrap(module.run_raptor)("tenant-1", "kb-1"))
res = inspect.unwrap(module.trace_index)("tenant-1", "kb-1")
assert res["code"] == module.RetCode.DATA_ERROR, res
# Invalid dataset ID
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True)
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None))
res = _run(inspect.unwrap(module.run_raptor)("tenant-1", "kb-1"))
res = inspect.unwrap(module.trace_index)("tenant-1", "kb-1")
assert "Invalid Dataset ID" in res["message"], res
stale_kb = _KB(kb_id="kb-1", raptor_task_id="task-old")
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, stale_kb))
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (False, None))
monkeypatch.setattr(module.DocumentService, "get_by_kb_id", lambda **_kwargs: ([{"id": "doc-1"}], 1))
monkeypatch.setattr(module.dataset_api_service, "queue_raptor_o_graphrag_tasks", lambda **_kwargs: "task-new")
monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: True)
res = _run(inspect.unwrap(module.run_raptor)("tenant-1", "kb-1"))
# No existing task — returns empty
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _KB(kb_id="kb-1", graphrag_task_id="")))
res = inspect.unwrap(module.trace_index)("tenant-1", "kb-1")
assert res["code"] == module.RetCode.SUCCESS, res
assert any("RAPTOR" in msg for msg in warnings), warnings
assert res["data"] == {}, res
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (True, SimpleNamespace(progress=0)))
res = _run(inspect.unwrap(module.run_raptor)("tenant-1", "kb-1"))
assert "already running" in res["message"], res
warnings.clear()
no_task_kb = _KB(kb_id="kb-1", raptor_task_id="")
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, no_task_kb))
monkeypatch.setattr(module.DocumentService, "get_by_kb_id", lambda **_kwargs: ([{"id": "doc-1"}], 1))
monkeypatch.setattr(module.dataset_api_service, "queue_raptor_o_graphrag_tasks", lambda **_kwargs: "queued-raptor")
monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: False)
res = _run(inspect.unwrap(module.run_raptor)("tenant-1", "kb-1"))
assert res["code"] == module.RetCode.SUCCESS, res
assert res["data"]["raptor_task_id"] == "queued-raptor", res
assert any("Cannot save raptor_task_id" in msg for msg in warnings), warnings
res = inspect.unwrap(module.trace_raptor)("tenant-1", "")
assert 'Dataset ID' in res["message"], res
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: False)
res = inspect.unwrap(module.trace_raptor)("tenant-1", "kb-1")
assert res["code"] == module.RetCode.DATA_ERROR, res
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda *_args, **_kwargs: True)
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None))
res = inspect.unwrap(module.trace_raptor)("tenant-1", "kb-1")
assert "Invalid Dataset ID" in res["message"], res
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _KB(kb_id="kb-1", raptor_task_id="task-1")))
# Task ID set but task not found — returns empty
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _KB(kb_id="kb-1", graphrag_task_id="task-1")))
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (False, None))
res = inspect.unwrap(module.trace_raptor)("tenant-1", "kb-1")
assert "RAPTOR Task Not Found" in res["message"], res
res = inspect.unwrap(module.trace_index)("tenant-1", "kb-1")
assert res["code"] == module.RetCode.SUCCESS, res
assert res["data"] == {}, res
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (True, SimpleNamespace(to_dict=lambda: {"id": _task_id, "progress": -1})))
res = inspect.unwrap(module.trace_raptor)("tenant-1", "kb-1")
# Task found — returns task data
monkeypatch.setattr(module.TaskService, "get_by_id", lambda _task_id: (True, SimpleNamespace(to_dict=lambda: {"id": _task_id, "progress": 1})))
res = inspect.unwrap(module.trace_index)("tenant-1", "kb-1")
assert res["code"] == module.RetCode.SUCCESS, res
assert res["data"]["id"] == "task-1", res

View File

@@ -1,662 +0,0 @@
#
# 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
from types import SimpleNamespace
import pytest
from test_common import (
delete_document,
document_change_status,
document_filter,
document_infos,
document_metadata_summary,
document_metadata_update,
document_update_metadata_setting,
)
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
INVALID_AUTH_CASES = [
(None, 401, "Unauthorized"),
(RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "Unauthorized"),
]
class TestAuthorization:
@pytest.mark.p2
@pytest.mark.parametrize("invalid_auth, expected_code, expected_fragment", INVALID_AUTH_CASES)
def test_filter_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = document_filter(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_infos_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = document_infos(invalid_auth, "kb_id", {"doc_ids": ["doc_id"]})
assert res["code"] == expected_code, res
assert expected_fragment in res["message"], res
## The inputs has been changed to add 'doc_ids'
## TODO:
#@pytest.mark.p2
#@pytest.mark.parametrize("invalid_auth, expected_code, expected_fragment", INVALID_AUTH_CASES)
#def test_metadata_summary_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
# res = document_metadata_summary(invalid_auth, {"kb_id": "kb_id"})
# assert res["code"] == expected_code, res
# assert expected_fragment in res["message"], res
## The inputs has been changed to deprecate 'selector'
## TODO:
#@pytest.mark.p2
#@pytest.mark.parametrize("invalid_auth, expected_code, expected_fragment", INVALID_AUTH_CASES)
#def test_metadata_update_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
# res = document_metadata_update(invalid_auth, {"kb_id": "kb_id", "selector": {"document_ids": ["doc_id"]}, "updates": []})
# 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_update_metadata_setting_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = document_update_metadata_setting(invalid_auth, "kb_id", "doc_id", {"metadata": {}})
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_change_status_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = document_change_status(invalid_auth, {"doc_ids": ["doc_id"], "status": "1"})
assert res["code"] == expected_code, res
assert expected_fragment in res["message"], res
class TestDocumentMetadata:
@pytest.mark.p2
def test_filter(self, WebApiAuth, add_dataset_func):
kb_id = add_dataset_func
res = document_filter(WebApiAuth, kb_id, {})
assert res["code"] == 0, res
assert "filter" in res["data"], res
assert "total" in res["data"], res
@pytest.mark.p2
def test_infos(self, WebApiAuth, add_document_func):
dataset_id, doc_id = add_document_func
res = document_infos(WebApiAuth, dataset_id, {"ids": [doc_id]})
assert res["code"] == 0, res
docs = res["data"]["docs"]
assert len(docs) == 1, docs
assert docs[0]["id"] == doc_id, res
## The inputs has been changed to add 'doc_ids'
## TODO:
#@pytest.mark.p2
#def test_metadata_summary(self, WebApiAuth, add_document_func):
# kb_id, _ = add_document_func
# res = document_metadata_summary(WebApiAuth, {"kb_id": kb_id})
# assert res["code"] == 0, res
# assert isinstance(res["data"]["summary"], dict), res
## The inputs has been changed to deprecate 'selector'
## TODO:
#@pytest.mark.p2
#def test_metadata_update(self, WebApiAuth, add_document_func):
# kb_id, doc_id = add_document_func
# payload = {
# "kb_id": kb_id,
# "selector": {"document_ids": [doc_id]},
# "updates": [{"key": "author", "value": "alice"}],
# "deletes": [],
# }
# res = document_metadata_update(WebApiAuth, payload)
# assert res["code"] == 0, res
# assert res["data"]["matched_docs"] == 1, res
# info_res = document_infos(WebApiAuth, {"doc_ids": [doc_id]})
# assert info_res["code"] == 0, info_res
# meta_fields = info_res["data"][0].get("meta_fields", {})
# assert meta_fields.get("author") == "alice", info_res
## The inputs has been changed to deprecate 'selector'
## TODO:
#@pytest.mark.p2
#def test_update_metadata_setting(self, WebApiAuth, add_document_func):
# _, doc_id = add_document_func
# metadata = {"source": "test"}
# res = document_update_metadata_setting(WebApiAuth, {"doc_id": doc_id, "metadata": metadata})
# assert res["code"] == 0, res
# assert res["data"]["id"] == doc_id, res
# assert res["data"]["parser_config"]["metadata"] == metadata, res
@pytest.mark.p2
def test_change_status(self, WebApiAuth, add_document_func):
dataset_id, doc_id = add_document_func
res = document_change_status(WebApiAuth, {"doc_ids": [doc_id], "status": "1"})
assert res["code"] == 0, res
assert res["data"][doc_id]["status"] == "1", res
info_res = document_infos(WebApiAuth, dataset_id, {"ids": [doc_id]})
assert info_res["code"] == 0, info_res
assert info_res["data"]["docs"][0]["status"] == "1", info_res
class TestDocumentMetadataNegative:
@pytest.mark.p2
def test_filter_missing_kb_id(self, WebApiAuth, add_document_func):
kb_id, doc_id = add_document_func
res = document_filter(WebApiAuth, "", {"ids": [doc_id]})
assert res["code"] == 100, res
assert "<MethodNotAllowed '405: Method Not Allowed'>" == res["message"], res
@pytest.mark.p3
def test_metadata_summary_missing_kb_id(self, WebApiAuth, add_document_func):
_, doc_id = add_document_func
res = document_metadata_summary(WebApiAuth, {"doc_ids": [doc_id]})
assert res["code"] == 101, res
assert "KB ID" in res["message"], res
## The inputs has been changed to deprecate 'selector'
## TODO:
#@pytest.mark.p3
#def test_metadata_update_missing_kb_id(self, WebApiAuth, add_document_func):
# _, doc_id = add_document_func
# res = document_metadata_update(WebApiAuth, {"selector": {"document_ids": [doc_id]}, "updates": []})
# assert res["code"] == 101, res
# assert "KB ID" in res["message"], res
@pytest.mark.p3
def test_infos_invalid_doc_id(self, WebApiAuth):
res = document_infos(WebApiAuth, {"doc_ids": ["invalid_id"]})
assert res["code"] == 109, res
assert "No authorization" in res["message"], res
@pytest.mark.p3
def test_update_metadata_setting_missing_metadata(self, WebApiAuth, add_document_func):
_, doc_id = add_document_func
res = document_update_metadata_setting(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 101, res
assert "required argument are missing" in res["message"], res
assert "metadata" in res["message"], res
@pytest.mark.p2
def test_update_metadata_setting_not_found(self, WebApiAuth, add_document_func):
"""Test updating metadata setting for a non-existent document returns error."""
dataset_id, doc_id = add_document_func
# First delete the document
delete_res = delete_document(WebApiAuth, dataset_id, {"ids": [doc_id]})
assert delete_res["code"] == 0, delete_res
# Now try to update metadata setting for the deleted document
res = document_update_metadata_setting(WebApiAuth, dataset_id, doc_id, {"metadata": {"author": "test"}})
assert res["code"] == 102, res
assert f"Document {doc_id} not found in dataset {dataset_id}" in res["message"], res
@pytest.mark.p3
def test_change_status_invalid_status(self, WebApiAuth, add_document_func):
_, doc_id = add_document_func
res = document_change_status(WebApiAuth, {"doc_ids": [doc_id], "status": "2"})
assert res["code"] == 101, res
assert "Status" in res["message"], res
def _run(coro):
return asyncio.run(coro)
class _DummyArgs:
def __init__(self, args=None):
self._args = args or {}
def get(self, key, default=None):
return self._args.get(key, default)
def getlist(self, key):
value = self._args.get(key, [])
if isinstance(value, list):
return value
return [value]
class _DummyRequest:
def __init__(self, args=None):
self.args = _DummyArgs(args)
class _DummyResponse:
def __init__(self, data=None):
self.data = data
self.headers = {}
@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)
@pytest.mark.p3
def test_update_metadata_missing_dataset_id(self, WebApiAuth, add_document_func):
"""Test the new unified update_metadata API - missing dataset_id."""
# Call with empty dataset_id (should fail validation)
res = document_metadata_update(WebApiAuth, "", {"dataset_id": "", "selector": {"document_ids": ["doc1"]}, "updates": []})
assert res["code"] == 404
assert res["message"] == "Not Found: /api/v1/datasets//documents/metadatas", res
@pytest.mark.p3
def test_update_metadata_success(self, WebApiAuth, add_document_func):
"""Test the new unified update_metadata API - success case."""
kb_id, doc_id = add_document_func
res = document_metadata_update(
WebApiAuth, kb_id,
{
"selector": {"document_ids": [doc_id]},
"updates": [{"key": "author", "value": "test_author"}],
"deletes": []
}
)
assert res["code"] == 0, res
@pytest.mark.p3
def test_update_metadata_invalid_delete_item(self, WebApiAuth, add_document_func):
"""Test the new unified update_metadata API - invalid delete item."""
kb_id, doc_id = add_document_func
res = document_metadata_update(
WebApiAuth, kb_id,
{
"selector": {"document_ids": [doc_id]},
"updates": [],
"deletes": [{}] # Invalid - missing key
}
)
assert res["code"] == 102
assert "Each delete requires key" in res["message"], res
def test_thumbnails_missing_ids_rewrite_and_exception_unit(self, document_app_module, monkeypatch):
module = document_app_module
monkeypatch.setattr(module, "request", _DummyRequest(args={}))
res = module.thumbnails()
assert res["code"] == module.RetCode.ARGUMENT_ERROR
assert 'Lack of "Document ID"' in res["message"]
monkeypatch.setattr(module, "request", _DummyRequest(args={"doc_ids": ["doc1", "doc2"]}))
monkeypatch.setattr(
module.DocumentService,
"get_thumbnails",
lambda _doc_ids: [
{"id": "doc1", "kb_id": "kb1", "thumbnail": "thumb.jpg"},
{"id": "doc2", "kb_id": "kb1", "thumbnail": f"{module.IMG_BASE64_PREFIX}blob"},
],
)
res = module.thumbnails()
assert res["code"] == 0
assert res["data"]["doc1"] == "/v1/document/image/kb1-thumb.jpg"
assert res["data"]["doc2"] == f"{module.IMG_BASE64_PREFIX}blob"
def raise_error(*_args, **_kwargs):
raise RuntimeError("thumb boom")
monkeypatch.setattr(module.DocumentService, "get_thumbnails", raise_error)
monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)})
res = module.thumbnails()
assert res["code"] == 500
assert "thumb boom" in res["message"]
def test_change_status_partial_failure_matrix_unit(self, document_app_module, monkeypatch):
module = document_app_module
calls = {"docstore_update": []}
doc_ids = ["unauth", "missing_doc", "missing_kb", "update_fail", "docstore_3022", "docstore_generic", "outer_exc"]
async def fake_request_json():
return {"doc_ids": doc_ids, "status": "1"}
def fake_accessible(doc_id, _uid):
return doc_id != "unauth"
def fake_get_by_id(doc_id):
if doc_id == "missing_doc":
return False, None
if doc_id == "outer_exc":
raise RuntimeError("explode")
kb_id = "kb_missing" if doc_id == "missing_kb" else "kb1"
chunk_num = 1 if doc_id in {"docstore_3022", "docstore_generic"} else 0
doc = SimpleNamespace(id=doc_id, kb_id=kb_id, status="0", chunk_num=chunk_num)
return True, doc
def fake_get_kb(kb_id):
if kb_id == "kb_missing":
return False, None
return True, SimpleNamespace(tenant_id="tenant1")
def fake_update_by_id(doc_id, _payload):
return doc_id != "update_fail"
class _DocStore:
def update(self, where, _payload, _index_name, _kb_id):
calls["docstore_update"].append(where["doc_id"])
if where["doc_id"] == "docstore_3022":
raise RuntimeError("3022 table missing")
if where["doc_id"] == "docstore_generic":
raise RuntimeError("doc store down")
return True
monkeypatch.setattr(module, "get_request_json", fake_request_json)
monkeypatch.setattr(module.DocumentService, "accessible", fake_accessible)
monkeypatch.setattr(module.DocumentService, "get_by_id", fake_get_by_id)
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda kb_id: fake_get_kb(kb_id))
monkeypatch.setattr(module.DocumentService, "update_by_id", fake_update_by_id)
monkeypatch.setattr(module.settings, "docStoreConn", _DocStore())
monkeypatch.setattr(module.search, "index_name", lambda tenant_id: f"idx_{tenant_id}")
res = _run(module.change_status.__wrapped__())
assert res["code"] == module.RetCode.SERVER_ERROR
assert res["message"] == "Partial failure"
assert res["data"]["unauth"]["error"] == "No authorization."
assert res["data"]["missing_doc"]["error"] == "No authorization."
assert res["data"]["missing_kb"]["error"] == "Can't find this dataset!"
assert res["data"]["update_fail"]["error"] == "Database error (Document update)!"
assert res["data"]["docstore_3022"]["error"] == "Document store table missing."
assert "Document store update failed:" in res["data"]["docstore_generic"]["error"]
assert "Internal server error: explode" == res["data"]["outer_exc"]["error"]
assert calls["docstore_update"] == ["docstore_3022", "docstore_generic"]
def test_change_status_invalid_status_unit(self, document_app_module, monkeypatch):
module = document_app_module
async def fake_request_json():
return {"doc_ids": ["doc1"], "status": "2"}
monkeypatch.setattr(module, "get_request_json", fake_request_json)
res = _run(module.change_status.__wrapped__())
assert res["code"] == module.RetCode.ARGUMENT_ERROR
assert '"Status" must be either 0 or 1!' in res["message"]
def test_change_status_all_success_unit(self, document_app_module, monkeypatch):
module = document_app_module
async def fake_request_json():
return {"doc_ids": ["doc1"], "status": "1"}
monkeypatch.setattr(module, "get_request_json", fake_request_json)
monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, SimpleNamespace(id="doc1", kb_id="kb1", status="0", chunk_num=0)))
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant1")))
monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: True)
res = _run(module.change_status.__wrapped__())
assert res["code"] == 0
assert res["data"]["doc1"]["status"] == "1"
def test_get_route_not_found_success_and_exception_unit(self, document_app_module, monkeypatch):
module = document_app_module
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None))
res = _run(module.get("doc1"))
assert res["code"] == module.RetCode.DATA_ERROR
assert "Document not found!" in res["message"]
async def fake_thread_pool_exec(*_args, **_kwargs):
return b"blob-data"
async def fake_make_response(data):
return _DummyResponse(data)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, SimpleNamespace(name="image.abc", type=module.FileType.VISUAL.value)))
monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("bucket", "name"))
monkeypatch.setattr(module.settings, "STORAGE_IMPL", SimpleNamespace(get=lambda *_args, **_kwargs: b"blob-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.get("doc1"))
assert isinstance(res, _DummyResponse)
assert res.data == b"blob-data"
assert res.headers["content_type"] == "image/abc"
assert res.headers["extension"] == "abc"
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (_ for _ in ()).throw(RuntimeError("get boom")))
monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)})
res = _run(module.get("doc1"))
assert res["code"] == 500
assert "get boom" in res["message"]
def test_download_attachment_success_and_exception_unit(self, document_app_module, monkeypatch):
module = document_app_module
monkeypatch.setattr(module, "request", _DummyRequest(args={"ext": "abc"}))
async def fake_thread_pool_exec(*_args, **_kwargs):
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.settings, "STORAGE_IMPL", SimpleNamespace(get=lambda *_args, **_kwargs: b"attachment"))
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("att1"))
assert isinstance(res, _DummyResponse)
assert res.data == b"attachment"
assert res.headers["content_type"] == "application/abc"
assert res.headers["extension"] == "abc"
async def raise_error(*_args, **_kwargs):
raise RuntimeError("download boom")
monkeypatch.setattr(module, "thread_pool_exec", raise_error)
monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)})
res = _run(module.download_attachment("att1"))
assert res["code"] == 500
assert "download boom" in res["message"]
def test_change_parser_guards_and_reset_update_failure_unit(self, document_app_module, monkeypatch):
module = document_app_module
monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)})
async def req_auth_fail():
return {"doc_id": "doc1", "parser_id": "naive", "pipeline_id": "pipe2"}
monkeypatch.setattr(module, "get_request_json", req_auth_fail)
monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: False)
res = _run(module.change_parser.__wrapped__())
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR
monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None))
res = _run(module.change_parser.__wrapped__())
assert res["code"] == module.RetCode.DATA_ERROR
assert "Document not found!" in res["message"]
async def req_same_pipeline():
return {"doc_id": "doc1", "parser_id": "naive", "pipeline_id": "pipe1"}
doc_same = SimpleNamespace(
id="doc1",
pipeline_id="pipe1",
parser_id="naive",
parser_config={"k": "v"},
token_num=0,
chunk_num=0,
process_duration=0,
kb_id="kb1",
type="doc",
name="doc.txt",
)
monkeypatch.setattr(module, "get_request_json", req_same_pipeline)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, doc_same))
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 0
calls = []
async def req_pipeline_change():
return {"doc_id": "doc1", "parser_id": "naive", "pipeline_id": "pipe2"}
doc = SimpleNamespace(
id="doc1",
pipeline_id="pipe1",
parser_id="naive",
parser_config={},
token_num=0,
chunk_num=0,
process_duration=0,
kb_id="kb1",
type="doc",
name="doc.txt",
)
def fake_update_by_id(doc_id, payload):
calls.append((doc_id, payload))
return True
monkeypatch.setattr(module, "get_request_json", req_pipeline_change)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, doc))
monkeypatch.setattr(module.DocumentService, "update_by_id", fake_update_by_id)
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 0
assert calls[0][1] == {"pipeline_id": "pipe2"}
assert calls[1][1]["run"] == module.TaskStatus.UNSTART.value
doc.token_num = 3
doc.chunk_num = 2
doc.process_duration = 9
monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *_args, **_kwargs: False)
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 0
monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *_args, **_kwargs: True)
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: None)
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 0
side_effects = {"img": [], "delete": []}
class _DocStore:
def index_exist(self, _idx, _kb_id):
return True
def delete(self, where, _idx, kb_id):
side_effects["delete"].append((where["doc_id"], kb_id))
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant1")
monkeypatch.setattr(module.DocumentService, "delete_chunk_images", lambda _doc, _tenant: side_effects["img"].append((_doc.id, _tenant)))
monkeypatch.setattr(module.search, "index_name", lambda tenant_id: f"idx_{tenant_id}")
monkeypatch.setattr(module.settings, "docStoreConn", _DocStore())
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 0
assert ("doc1", "tenant1") in side_effects["img"]
assert ("doc1", "kb1") in side_effects["delete"]
async def req_same_parser_with_cfg():
return {"doc_id": "doc1", "parser_id": "naive", "parser_config": {"a": 1}}
doc_same_parser = SimpleNamespace(
id="doc1",
pipeline_id="pipe1",
parser_id="naive",
parser_config={"a": 1},
token_num=0,
chunk_num=0,
process_duration=0,
kb_id="kb1",
type="doc",
name="doc.txt",
)
monkeypatch.setattr(module, "get_request_json", req_same_parser_with_cfg)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, doc_same_parser))
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 0
async def req_same_parser_no_cfg():
return {"doc_id": "doc1", "parser_id": "naive"}
monkeypatch.setattr(module, "get_request_json", req_same_parser_no_cfg)
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 0
parser_cfg_updates = []
async def req_parser_update():
return {"doc_id": "doc1", "parser_id": "paper", "pipeline_id": "", "parser_config": {"beta": True}}
doc_parser_update = SimpleNamespace(
id="doc1",
pipeline_id="pipe1",
parser_id="naive",
parser_config={"alpha": 1},
token_num=0,
chunk_num=0,
process_duration=0,
kb_id="kb1",
type="doc",
name="doc.txt",
)
monkeypatch.setattr(module, "get_request_json", req_parser_update)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, doc_parser_update))
monkeypatch.setattr(module.DocumentService, "update_parser_config", lambda doc_id, cfg: parser_cfg_updates.append((doc_id, cfg)))
monkeypatch.setattr(module.DocumentService, "update_by_id", lambda *_args, **_kwargs: True)
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 0
assert parser_cfg_updates == [("doc1", {"beta": True})]
def raise_parser_config(*_args, **_kwargs):
raise RuntimeError("parser boom")
monkeypatch.setattr(module.DocumentService, "update_parser_config", raise_parser_config)
res = _run(module.change_parser.__wrapped__())
assert res["code"] == 500
assert "parser boom" in res["message"]
def test_get_image_success_and_exception_unit(self, document_app_module, monkeypatch):
module = document_app_module
class _Headers(dict):
def set(self, key, value):
self[key] = value
class _ImageResponse:
def __init__(self, data):
self.data = data
self.headers = _Headers()
async def fake_thread_pool_exec(*_args, **_kwargs):
return b"image-bytes"
async def fake_make_response(data):
return _ImageResponse(data)
monkeypatch.setattr(module, "thread_pool_exec", fake_thread_pool_exec)
monkeypatch.setattr(module, "make_response", fake_make_response)
monkeypatch.setattr(module.settings, "STORAGE_IMPL", SimpleNamespace(get=lambda *_args, **_kwargs: b"image-bytes"))
res = _run(module.get_image("bucket-name"))
assert isinstance(res, _ImageResponse)
assert res.data == b"image-bytes"
assert res.headers["Content-Type"] == "image/JPEG"
async def raise_error(*_args, **_kwargs):
raise RuntimeError("image boom")
monkeypatch.setattr(module, "thread_pool_exec", raise_error)
monkeypatch.setattr(module, "server_error_response", lambda e: {"code": 500, "message": str(e)})
res = _run(module.get_image("bucket-name"))
assert res["code"] == 500
assert "image boom" in res["message"]

View File

@@ -172,15 +172,15 @@ class TestDocumentsList:
def test_missing_kb_id(self, WebApiAuth):
"""Test missing KB ID returns error."""
res = list_documents(WebApiAuth, {"kb_id": ""})
assert res["code"] == 100
assert res["message"] == "<MethodNotAllowed '405: Method Not Allowed'>"
assert res["code"] == 102
assert res["message"]
@pytest.mark.p2
def test_unauthorized_dataset(self, WebApiAuth):
"""Test unauthorized dataset returns error."""
res = list_documents(WebApiAuth, {"kb_id": "non_existent_kb_id"})
assert res["code"] == 102
assert "You don't own the dataset" in res["message"]
assert res["message"]
@pytest.mark.p3
def test_invalid_run_status_filter(self, WebApiAuth, add_documents):

View File

@@ -1,50 +0,0 @@
#
# 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 pytest
from test_common import batch_create_datasets, list_datasets, delete_datasets
from libs.auth import RAGFlowWebApiAuth
from pytest import FixtureRequest
from ragflow_sdk import RAGFlow
@pytest.fixture(scope="class")
def add_datasets(request: FixtureRequest, client: RAGFlow, WebApiAuth: RAGFlowWebApiAuth) -> list[str]:
dataset_ids = batch_create_datasets(WebApiAuth, 5)
def cleanup():
# Web KB cleanup cannot call SDK dataset bulk delete with empty ids; deletion must stay explicit.
res = list_datasets(WebApiAuth, params={"page_size": 1000})
existing_ids = {kb["id"] for kb in res["data"]}
ids_to_delete = list({dataset_id for dataset_id in dataset_ids if dataset_id in existing_ids})
delete_datasets(WebApiAuth, {"ids": ids_to_delete})
request.addfinalizer(cleanup)
return dataset_ids
@pytest.fixture(scope="function")
def add_datasets_func(request: FixtureRequest, client: RAGFlow, WebApiAuth: RAGFlowWebApiAuth) -> list[str]:
dataset_ids = batch_create_datasets(WebApiAuth, 3)
def cleanup():
# Web KB cleanup cannot call SDK dataset bulk delete with empty ids; deletion must stay explicit.
res = list_datasets(WebApiAuth, params={"page_size": 1000})
existing_ids = {kb["id"] for kb in res["data"]}
ids_to_delete = list({dataset_id for dataset_id in dataset_ids if dataset_id in existing_ids})
delete_datasets(WebApiAuth, {"ids": ids_to_delete})
request.addfinalizer(cleanup)
return dataset_ids

View File

@@ -1,109 +0,0 @@
#
# 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.
#
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
from test_common import create_dataset
from configs import DATASET_NAME_LIMIT, INVALID_API_TOKEN
from hypothesis import example, given, settings
from libs.auth import RAGFlowWebApiAuth
from utils.hypothesis_utils import valid_names
@pytest.mark.usefixtures("clear_datasets")
class TestAuthorization:
@pytest.mark.p2
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
ids=["empty_auth", "invalid_api_token"],
)
def test_auth_invalid(self, invalid_auth, expected_code, expected_message):
res = create_dataset(invalid_auth, {"name": "auth_test"})
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
@pytest.mark.usefixtures("clear_datasets")
class TestCapability:
@pytest.mark.p3
def test_create_kb_1k(self, WebApiAuth):
for i in range(1_000):
payload = {"name": f"dataset_{i}"}
res = create_dataset(WebApiAuth, payload)
assert res["code"] == 0, f"Failed to create dataset {i}"
@pytest.mark.p3
def test_create_kb_concurrent(self, WebApiAuth):
count = 100
with ThreadPoolExecutor(max_workers=5) as executor:
futures = [executor.submit(create_dataset, WebApiAuth, {"name": f"dataset_{i}"}) for i in range(count)]
responses = list(as_completed(futures))
assert len(responses) == count, responses
assert all(future.result()["code"] == 0 for future in futures)
@pytest.mark.usefixtures("clear_datasets")
class TestDatasetCreate:
@pytest.mark.p1
@given(name=valid_names())
@example("a" * 128)
@settings(max_examples=20)
def test_name(self, WebApiAuth, name):
res = create_dataset(WebApiAuth, {"name": name})
assert res["code"] == 0, res
@pytest.mark.p2
@pytest.mark.parametrize(
"name, expected_message",
[
("", "Field: <name> - Message: <String should have at least 1 character>"),
(" ", "Field: <name> - Message: <String should have at least 1 character>"),
("a" * (DATASET_NAME_LIMIT + 1), "Field: <name> - Message: <String should have at most 128 characters>"),
(0, "Field: <name> - Message: <Input should be a valid string>"),
(None, "Field: <name> - Message: <Input should be a valid string>"),
],
ids=["empty_name", "space_name", "too_long_name", "invalid_name", "None_name"],
)
def test_name_invalid(self, WebApiAuth, name, expected_message):
payload = {"name": name}
res = create_dataset(WebApiAuth, payload)
assert res["code"] == 101, res
assert expected_message in res["message"], res
@pytest.mark.p3
def test_name_duplicated(self, WebApiAuth):
name = "duplicated_name"
payload = {"name": name}
res = create_dataset(WebApiAuth, payload)
assert res["code"] == 0, res
res = create_dataset(WebApiAuth, payload)
assert res["code"] == 0, res
@pytest.mark.p3
def test_name_case_insensitive(self, WebApiAuth):
name = "CaseInsensitive"
payload = {"name": name.upper()}
res = create_dataset(WebApiAuth, payload)
assert res["code"] == 0, res
payload = {"name": name.lower()}
res = create_dataset(WebApiAuth, payload)
assert res["code"] == 0, res

View File

@@ -1,53 +0,0 @@
#
# 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 pytest
from test_common import (
detail_kb,
)
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
class TestAuthorization:
@pytest.mark.p2
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
)
def test_auth_invalid(self, invalid_auth, expected_code, expected_message):
res = detail_kb(invalid_auth)
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
class TestDatasetsDetail:
@pytest.mark.p1
def test_kb_id(self, WebApiAuth, add_dataset):
kb_id = add_dataset
payload = {"kb_id": kb_id}
res = detail_kb(WebApiAuth, payload)
assert res["code"] == 0, res
assert res["data"]["name"] == "kb_0"
@pytest.mark.p2
def test_id_wrong_uuid(self, WebApiAuth):
payload = {"kb_id": "d94a8dc02c9711f0930f7fbc369eab6d"}
res = detail_kb(WebApiAuth, payload)
assert res["code"] == 103, res
assert "Only owner of dataset authorized for this operation." in res["message"], res

View File

@@ -1,233 +0,0 @@
#
# 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 pytest
from test_common import (
kb_delete_pipeline_logs,
kb_list_pipeline_dataset_logs,
kb_list_pipeline_logs,
kb_pipeline_log_detail,
run_graphrag,
trace_graphrag,
run_raptor,
trace_raptor,
kb_run_mindmap,
kb_trace_mindmap,
list_documents,
parse_documents,
)
from utils import wait_for
TASK_STATUS_DONE = "3"
def _find_task(data, task_id):
if isinstance(data, dict):
if data.get("id") == task_id:
return data
tasks = data.get("tasks")
if isinstance(tasks, list):
for item in tasks:
if isinstance(item, dict) and item.get("id") == task_id:
return item
elif isinstance(data, list):
for item in data:
if isinstance(item, dict) and item.get("id") == task_id:
return item
return None
def _assert_progress_in_scale(progress, payload):
assert isinstance(progress, (int, float)), payload
if progress < 0:
assert False, f"Negative progress is not expected: {payload}"
scale = 100 if progress > 1 else 1
# Infer scale from observed payload (0..1 or 0..100).
assert 0 <= progress <= scale, payload
return scale
def _wait_for_task(trace_func, auth, kb_id, task_id, timeout=60, use_params_payload=False):
@wait_for(timeout, 1, "Pipeline task trace timeout")
def _condition():
if use_params_payload:
res = trace_func(auth, {"kb_id": kb_id})
else:
res = trace_func(auth, kb_id)
if res["code"] != 0:
return False
return _find_task(res["data"], task_id) is not None
_condition()
def _wait_for_docs_parsed(auth, kb_id, timeout=60):
@wait_for(timeout, 2, "Document parsing timeout")
def _condition():
res = list_documents(auth, {"kb_id": kb_id})
if res["code"] != 0:
return False
for doc in res["data"]["docs"]:
progress = doc.get("progress", 0)
_assert_progress_in_scale(progress, doc)
scale = 100 if progress > 1 else 1
if doc.get("run") != TASK_STATUS_DONE or progress < scale:
return False
return True
_condition()
def _wait_for_pipeline_logs(auth, kb_id, timeout=30):
@wait_for(timeout, 1, "Pipeline log timeout")
def _condition():
res = kb_list_pipeline_logs(auth, params={"kb_id": kb_id}, payload={})
if res["code"] != 0:
return False
return bool(res["data"]["logs"])
_condition()
class TestKbPipelineTasks:
@pytest.mark.p3
def test_graphrag_run_and_trace(self, WebApiAuth, add_chunks):
kb_id, _, _ = add_chunks
run_res = run_graphrag(WebApiAuth, kb_id)
assert run_res["code"] == 0, run_res
task_id = run_res["data"]["graphrag_task_id"]
assert task_id, run_res
_wait_for_task(trace_graphrag, WebApiAuth, kb_id, task_id)
trace_res = trace_graphrag(WebApiAuth, kb_id)
assert trace_res["code"] == 0, trace_res
task = _find_task(trace_res["data"], task_id)
assert task, trace_res
assert task["id"] == task_id, trace_res
progress = task.get("progress")
_assert_progress_in_scale(progress, task)
@pytest.mark.p3
def test_raptor_run_and_trace(self, WebApiAuth, add_chunks):
kb_id, _, _ = add_chunks
run_res = run_raptor(WebApiAuth, kb_id)
assert run_res["code"] == 0, run_res
task_id = run_res["data"]["raptor_task_id"]
assert task_id, run_res
_wait_for_task(trace_raptor, WebApiAuth, kb_id, task_id)
trace_res = trace_raptor(WebApiAuth, kb_id)
assert trace_res["code"] == 0, trace_res
task = _find_task(trace_res["data"], task_id)
assert task, trace_res
assert task["id"] == task_id, trace_res
progress = task.get("progress")
_assert_progress_in_scale(progress, task)
@pytest.mark.p3
def test_mindmap_run_and_trace(self, WebApiAuth, add_chunks):
kb_id, _, _ = add_chunks
run_res = kb_run_mindmap(WebApiAuth, {"kb_id": kb_id})
assert run_res["code"] == 0, run_res
task_id = run_res["data"]["mindmap_task_id"]
assert task_id, run_res
_wait_for_task(kb_trace_mindmap, WebApiAuth, kb_id, task_id, use_params_payload=True)
trace_res = kb_trace_mindmap(WebApiAuth, {"kb_id": kb_id})
assert trace_res["code"] == 0, trace_res
task = _find_task(trace_res["data"], task_id)
assert task, trace_res
assert task["id"] == task_id, trace_res
progress = task.get("progress")
_assert_progress_in_scale(progress, task)
class TestKbPipelineLogs:
@pytest.mark.p3
def test_pipeline_log_lifecycle(self, WebApiAuth, add_document):
kb_id, document_id = add_document
parse_documents(WebApiAuth, {"doc_ids": [document_id], "run": "1"})
_wait_for_docs_parsed(WebApiAuth, kb_id)
_wait_for_pipeline_logs(WebApiAuth, kb_id)
list_res = kb_list_pipeline_logs(WebApiAuth, params={"kb_id": kb_id}, payload={})
assert list_res["code"] == 0, list_res
assert "total" in list_res["data"], list_res
assert isinstance(list_res["data"]["logs"], list), list_res
assert list_res["data"]["logs"], list_res
log_id = list_res["data"]["logs"][0]["id"]
detail_res = kb_pipeline_log_detail(WebApiAuth, {"log_id": log_id})
assert detail_res["code"] == 0, detail_res
detail = detail_res["data"]
assert detail["id"] == log_id, detail_res
assert detail["kb_id"] == kb_id, detail_res
for key in ["document_id", "task_type", "operation_status", "progress"]:
assert key in detail, detail_res
delete_res = kb_delete_pipeline_logs(WebApiAuth, params={"kb_id": kb_id}, payload={"log_ids": [log_id]})
assert delete_res["code"] == 0, delete_res
assert delete_res["data"] is True, delete_res
@wait_for(30, 1, "Pipeline log delete timeout")
def _condition():
res = kb_list_pipeline_logs(WebApiAuth, params={"kb_id": kb_id}, payload={})
if res["code"] != 0:
return False
return all(log.get("id") != log_id for log in res["data"]["logs"])
_condition()
@pytest.mark.p3
def test_list_pipeline_dataset_logs(self, WebApiAuth, add_document):
kb_id, _ = add_document
res = kb_list_pipeline_dataset_logs(WebApiAuth, params={"kb_id": kb_id}, payload={})
assert res["code"] == 0, res
assert "total" in res["data"], res
assert isinstance(res["data"]["logs"], list), res
@pytest.mark.p3
def test_pipeline_log_detail_missing_id(self, WebApiAuth):
res = kb_pipeline_log_detail(WebApiAuth, {})
assert res["code"] == 101, res
assert "Pipeline log ID" in res["message"], res
@pytest.mark.p3
def test_delete_pipeline_logs_empty(self, WebApiAuth, add_document):
kb_id, _ = add_document
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

File diff suppressed because it is too large Load Diff

View File

@@ -1,296 +0,0 @@
#
# 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 uuid
import pytest
from test_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,
rm_tags,
update_chunk,
)
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
from utils import wait_for
INVALID_AUTH_CASES = [
(None, 401, "Unauthorized"),
(RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "Unauthorized"),
]
TAG_SEED_TIMEOUT = 20
def _wait_for_tag(auth, kb_id, tag, timeout=TAG_SEED_TIMEOUT):
@wait_for(timeout, 1, "Tag seed timeout")
def _condition():
res = list_tags(auth, kb_id)
if res["code"] != 0:
return False
return tag in res["data"]
try:
_condition()
except AssertionError:
return False
return True
def _seed_tag(auth, kb_id, document_id, chunk_id):
# KB tags are derived from chunk tag_kwd, not document metadata.
tag = f"tag_{uuid.uuid4().hex[:8]}"
res = update_chunk(
auth,
kb_id,
document_id,
chunk_id,
{
"content": f"tag seed {tag}",
"tag_kwd": [tag],
},
)
assert res["code"] == 0, res
if not _wait_for_tag(auth, kb_id, tag):
return None
return tag
class TestAuthorization:
@pytest.mark.p2
@pytest.mark.parametrize("invalid_auth, expected_code, expected_fragment", INVALID_AUTH_CASES)
def test_list_tags_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = list_tags(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_list_tags_from_kbs_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = list_tags_from_kbs(invalid_auth, {"kb_ids": "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_rm_tags_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = rm_tags(invalid_auth, "kb_id", {"tags": ["tag"]})
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_rename_tag_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = rename_tags(invalid_auth, "kb_id", {"from_tag": "old", "to_tag": "new"})
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_get_meta_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = kb_get_meta(invalid_auth, {"kb_ids": "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_basic_info_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = kb_basic_info(invalid_auth, {"kb_id": "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_update_metadata_setting_auth_invalid(self, invalid_auth, expected_code, expected_fragment):
res = kb_update_metadata_setting(invalid_auth, {"kb_id": "kb_id", "metadata": {}})
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
def test_list_tags(self, WebApiAuth, add_dataset):
kb_id = add_dataset
res = list_tags(WebApiAuth, kb_id)
assert res["code"] == 0, res
assert isinstance(res["data"], list), res
@pytest.mark.p2
def test_list_tags_from_kbs(self, WebApiAuth, add_dataset):
kb_id = add_dataset
res = list_tags_from_kbs(WebApiAuth, {"kb_ids": kb_id})
assert res["code"] == 0, res
assert isinstance(res["data"], list), res
@pytest.mark.p3
def test_rm_tags(self, WebApiAuth, add_chunks):
kb_id, document_id, chunk_ids = add_chunks
tag_to_remove = _seed_tag(WebApiAuth, kb_id, document_id, chunk_ids[0])
if not tag_to_remove:
# Tag aggregation is index-backed; skip if it never surfaces.
pytest.skip("Seeded tag did not appear in list_tags.")
res = rm_tags(WebApiAuth, kb_id, {"tags": [tag_to_remove]})
assert res["code"] == 0, res
assert res["data"] is True, res
@wait_for(TAG_SEED_TIMEOUT, 1, "Tag removal timeout")
def _condition():
after_res = list_tags(WebApiAuth, kb_id)
if after_res["code"] != 0:
return False
return tag_to_remove not in after_res["data"]
_condition()
@pytest.mark.p3
def test_rename_tag(self, WebApiAuth, add_chunks):
kb_id, document_id, chunk_ids = add_chunks
from_tag = _seed_tag(WebApiAuth, kb_id, document_id, chunk_ids[0])
if not from_tag:
# Tag aggregation is index-backed; skip if it never surfaces.
pytest.skip("Seeded tag did not appear in list_tags.")
to_tag = f"{from_tag}_renamed"
res = rename_tags(WebApiAuth, kb_id, {"from_tag": from_tag, "to_tag": to_tag})
assert res["code"] == 0, res
assert res["data"] is True, res
@wait_for(TAG_SEED_TIMEOUT, 1, "Tag rename timeout")
def _condition():
after_res = list_tags(WebApiAuth, kb_id)
if after_res["code"] != 0:
return False
tags = after_res["data"]
return to_tag in tags and from_tag not in tags
_condition()
@pytest.mark.p2
def test_get_meta(self, WebApiAuth, add_dataset):
kb_id = add_dataset
res = kb_get_meta(WebApiAuth, {"kb_ids": kb_id})
assert res["code"] == 0, res
assert isinstance(res["data"], dict), res
@pytest.mark.p2
def test_basic_info(self, WebApiAuth, add_dataset):
kb_id = add_dataset
res = kb_basic_info(WebApiAuth, {"kb_id": kb_id})
assert res["code"] == 0, res
for key in ["processing", "finished", "failed", "cancelled", "downloaded"]:
assert key in res["data"], res
@pytest.mark.p2
def test_update_metadata_setting(self, WebApiAuth, add_dataset):
kb_id = add_dataset
metadata = {"source": "test"}
res = kb_update_metadata_setting(WebApiAuth, {"kb_id": kb_id, "metadata": metadata, "enable_metadata": True})
assert res["code"] == 0, res
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
def test_list_tags_invalid_kb(self, WebApiAuth):
res = list_tags(WebApiAuth, "invalid_kb_id")
assert res["code"] == 109, res
assert "No authorization" in res["message"], res
@pytest.mark.p3
def test_list_tags_from_kbs_invalid_kb(self, WebApiAuth):
res = list_tags_from_kbs(WebApiAuth, {"kb_ids": "invalid_kb_id"})
assert res["code"] == 109, res
assert "No authorization" in res["message"], res
@pytest.mark.p3
def test_rm_tags_invalid_kb(self, WebApiAuth):
res = rm_tags(WebApiAuth, "invalid_kb_id", {"tags": ["tag"]})
assert res["code"] == 109, res
assert "No authorization" in res["message"], res
@pytest.mark.p3
def test_rename_tag_invalid_kb(self, WebApiAuth):
res = rename_tags(WebApiAuth, "invalid_kb_id", {"from_tag": "old", "to_tag": "new"})
assert res["code"] == 109, res
assert "No authorization" in res["message"], res
@pytest.mark.p3
def test_get_meta_invalid_kb(self, WebApiAuth):
res = kb_get_meta(WebApiAuth, {"kb_ids": "invalid_kb_id"})
assert res["code"] == 109, res
assert "No authorization" in res["message"], res
@pytest.mark.p3
def test_basic_info_invalid_kb(self, WebApiAuth):
res = kb_basic_info(WebApiAuth, {"kb_id": "invalid_kb_id"})
assert res["code"] == 109, res
assert "No authorization" in res["message"], res
@pytest.mark.p3
def test_update_metadata_setting_missing_metadata(self, WebApiAuth, add_dataset):
res = kb_update_metadata_setting(WebApiAuth, {"kb_id": add_dataset})
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

View File

@@ -1,201 +0,0 @@
#
# 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 json
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
from test_common import list_datasets
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
from utils import is_sorted
class TestAuthorization:
@pytest.mark.p2
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
)
def test_auth_invalid(self, invalid_auth, expected_code, expected_message):
res = list_datasets(invalid_auth)
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
class TestCapability:
@pytest.mark.p3
def test_concurrent_list(self, WebApiAuth):
count = 100
with ThreadPoolExecutor(max_workers=5) as executor:
futures = [executor.submit(list_datasets, WebApiAuth) for i in range(count)]
responses = list(as_completed(futures))
assert len(responses) == count, responses
assert all(future.result()["code"] == 0 for future in futures)
@pytest.mark.usefixtures("add_datasets")
class TestDatasetsList:
@pytest.mark.p2
def test_params_unset(self, WebApiAuth):
res = list_datasets(WebApiAuth, None)
assert res["code"] == 0, res
assert len(res["data"]) == 5, res
@pytest.mark.p2
def test_params_empty(self, WebApiAuth):
res = list_datasets(WebApiAuth, {})
assert res["code"] == 0, res
assert len(res["data"]) == 5, res
@pytest.mark.p1
@pytest.mark.parametrize(
"params, expected_page_size",
[
({"page": 2, "page_size": 2}, 2),
({"page": 3, "page_size": 2}, 1),
({"page": 4, "page_size": 2}, 0),
({"page": "2", "page_size": 2}, 2),
({"page": 1, "page_size": 10}, 5),
],
ids=["normal_middle_page", "normal_last_partial_page", "beyond_max_page", "string_page_number", "full_data_single_page"],
)
def test_page(self, WebApiAuth, params, expected_page_size):
res = list_datasets(WebApiAuth, params)
assert res["code"] == 0, res
assert len(res["data"]) == expected_page_size, res
@pytest.mark.skip
@pytest.mark.p2
@pytest.mark.parametrize(
"params, expected_code, expected_message",
[
({"page": 0}, 101, "Input should be greater than or equal to 1"),
({"page": "a"}, 101, "Input should be a valid integer, unable to parse string as an integer"),
],
ids=["page_0", "page_a"],
)
def test_page_invalid(self, WebApiAuth, params, expected_code, expected_message):
res = list_datasets(WebApiAuth, params=params)
assert res["code"] == expected_code, res
assert expected_message in res["message"], res
@pytest.mark.p2
def test_page_none(self, WebApiAuth):
params = {"page": None}
res = list_datasets(WebApiAuth, params)
assert res["code"] == 0, res
assert len(res["data"]) == 5, res
@pytest.mark.p1
@pytest.mark.parametrize(
"params, expected_page_size",
[
({"page": 1, "page_size": 1}, 1),
({"page": 1, "page_size": 3}, 3),
({"page": 1, "page_size": 5}, 5),
({"page": 1, "page_size": 6}, 5),
({"page": 1, "page_size": "1"}, 1),
],
ids=["min_valid_page_size", "medium_page_size", "page_size_equals_total", "page_size_exceeds_total", "string_type_page_size"],
)
def test_page_size(self, WebApiAuth, params, expected_page_size):
res = list_datasets(WebApiAuth, params)
assert res["code"] == 0, res
assert len(res["data"]) == expected_page_size, res
@pytest.mark.skip
@pytest.mark.p2
@pytest.mark.parametrize(
"params, expected_code, expected_message",
[
({"page_size": 0}, 101, "Input should be greater than or equal to 1"),
({"page_size": "a"}, 101, "Input should be a valid integer, unable to parse string as an integer"),
],
)
def test_page_size_invalid(self, WebApiAuth, params, expected_code, expected_message):
res = list_datasets(WebApiAuth, params)
assert res["code"] == expected_code, res
assert expected_message in res["message"], res
@pytest.mark.p2
def test_page_size_none(self, WebApiAuth):
params = {"page_size": None}
res = list_datasets(WebApiAuth, params)
assert res["code"] == 0, res
assert len(res["data"]) == 5, res
@pytest.mark.p3
@pytest.mark.parametrize(
"params, assertions",
[
({"orderby": "update_time"}, lambda r: (is_sorted(r["data"], "update_time", True))),
],
ids=["orderby_update_time"],
)
def test_orderby(self, WebApiAuth, params, assertions):
res = list_datasets(WebApiAuth, params)
assert res["code"] == 0, res
if callable(assertions):
assert assertions(res), res
@pytest.mark.p3
@pytest.mark.parametrize(
"params, assertions",
[
({"desc": "True"}, lambda r: (is_sorted(r["data"], "update_time", True))),
({"desc": "False"}, lambda r: (is_sorted(r["data"], "update_time", False))),
],
ids=["desc=True", "desc=False"],
)
def test_desc(self, WebApiAuth, params, assertions):
res = list_datasets(WebApiAuth, params)
assert res["code"] == 0, res
if callable(assertions):
assert assertions(res), res
@pytest.mark.p2
@pytest.mark.parametrize(
"params, expected_page_size",
[
({"ext": json.dumps({"parser_id": "naive"})}, 5),
({"ext": json.dumps({"parser_id": "qa"})}, 0),
],
ids=["naive", "dqa"],
)
def test_parser_id(self, WebApiAuth, params, expected_page_size):
res = list_datasets(WebApiAuth, params)
assert res["code"] == 0, res
assert len(res["data"]) == expected_page_size, res
@pytest.mark.p2
def test_owner_ids_payload_mode(self, WebApiAuth):
base_res = list_datasets(WebApiAuth, {"page_size": 10})
assert base_res["code"] == 0, base_res
assert base_res["data"], base_res
owner_id = base_res["data"][0]["tenant_id"]
res = list_datasets(
WebApiAuth,
params={"page": 1, "page_size": 2, "desc": "false", "ext": json.dumps({"owner_ids": [owner_id]})},
)
assert res["code"] == 0, res
assert res["total_datasets"] >= len(res["data"]), res
assert len(res["data"]) <= 2, res
assert all(kb["tenant_id"] == owner_id for kb in res["data"]), res

View File

@@ -1,61 +0,0 @@
#
# 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 pytest
from test_common import (
list_datasets,
delete_datasets,
)
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
class TestAuthorization:
@pytest.mark.p2
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
)
def test_auth_invalid(self, invalid_auth, expected_code, expected_message):
res = delete_datasets(invalid_auth)
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
class TestDatasetsDelete:
@pytest.mark.p1
def test_kb_id(self, WebApiAuth, add_datasets_func):
kb_ids = add_datasets_func
payload = {"ids": [kb_ids[0]]}
res = delete_datasets(WebApiAuth, payload)
assert res["code"] == 0, res
res = list_datasets(WebApiAuth)
assert len(res["data"]) == 2, res
@pytest.mark.p2
@pytest.mark.usefixtures("add_dataset_func")
def test_id_wrong_uuid(self, WebApiAuth):
payload = {"ids": ["d94a8dc02c9711f0930f7fbc369eab6d"]}
res = delete_datasets(WebApiAuth, payload)
assert res["code"] == 102, res
assert "lacks permission" in res["message"], res
res = list_datasets(WebApiAuth)
assert len(res["data"]) == 1, res

View File

@@ -1,382 +0,0 @@
#
# 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 os
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
from test_common import update_dataset
from configs import DATASET_NAME_LIMIT, INVALID_API_TOKEN
from hypothesis import HealthCheck, example, given, settings
from libs.auth import RAGFlowWebApiAuth
from utils import encode_avatar
from utils.file_utils import create_image_file
from utils.hypothesis_utils import valid_names
class TestAuthorization:
@pytest.mark.p2
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
ids=["empty_auth", "invalid_api_token"],
)
def test_auth_invalid(self, invalid_auth, expected_code, expected_message):
res = update_dataset(invalid_auth, "dataset_id")
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
class TestCapability:
@pytest.mark.p3
def test_update_dateset_concurrent(self, WebApiAuth, add_dataset_func):
dataset_id = add_dataset_func
count = 100
with ThreadPoolExecutor(max_workers=5) as executor:
futures = [
executor.submit(
update_dataset,
WebApiAuth,
dataset_id,
{
"name": f"dataset_{i}",
"description": "",
"chunk_method": "naive",
},
)
for i in range(count)
]
responses = list(as_completed(futures))
assert len(responses) == count, responses
assert all(future.result()["code"] == 0 for future in futures)
class TestDatasetUpdate:
@pytest.mark.p3
def test_dataset_id_not_uuid(self, WebApiAuth):
payload = {"name": "not uuid", "description": "", "chunk_method": "naive"}
res = update_dataset(WebApiAuth, "not_uuid", payload)
assert res["code"] == 101, res
assert "Invalid UUID1 format" in res["message"], res
@pytest.mark.p1
@given(name=valid_names())
@example("a" * 128)
# Network-bound API call; disable Hypothesis deadline to avoid flaky timeouts.
@settings(max_examples=20, suppress_health_check=[HealthCheck.function_scoped_fixture], deadline=None)
def test_name(self, WebApiAuth, add_dataset_func, name):
dataset_id = add_dataset_func
payload = {"name": name, "description": "", "chunk_method": "naive"}
res = update_dataset(WebApiAuth, dataset_id, payload)
assert res["code"] == 0, res
assert res["data"]["name"] == name, res
@pytest.mark.p2
@pytest.mark.parametrize(
"name, expected_message",
[
("", "Field: <name> - Message: <String should have at least 1 character>"),
(" ", "Field: <name> - Message: <String should have at least 1 character>"),
("a" * (DATASET_NAME_LIMIT + 1), "Field: <name> - Message: <String should have at most 128 characters>"),
(0, "Field: <name> - Message: <Input should be a valid string>"),
(None, "Field: <name> - Message: <Input should be a valid string>"),
],
ids=["empty_name", "space_name", "too_long_name", "invalid_name", "None_name"],
)
def test_name_invalid(self, WebApiAuth, add_dataset_func, name, expected_message):
kb_id = add_dataset_func
payload = {"name": name, "description": "", "chunk_method": "naive"}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 101, res
assert expected_message in res["message"], res
@pytest.mark.p3
def test_name_duplicated(self, WebApiAuth, add_datasets_func):
kb_id = add_datasets_func[0]
name = "kb_1"
payload = {"name": name, "description": "", "chunk_method": "naive"}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 102, res
assert res["message"] == "Dataset name 'kb_1' already exists", res
@pytest.mark.p3
def test_name_case_insensitive(self, WebApiAuth, add_datasets_func):
kb_id = add_datasets_func[0]
name = "KB_1"
payload = {"name": name, "description": "", "chunk_method": "naive"}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 102, res
assert res["message"] == "Dataset name 'KB_1' already exists", res
@pytest.mark.p2
def test_avatar(self, WebApiAuth, add_dataset_func, tmp_path):
kb_id = add_dataset_func
fn = create_image_file(tmp_path / "ragflow_test.png")
payload = {
"name": "avatar",
"description": "",
"chunk_method": "naive",
"avatar": f"data:image/png;base64,{encode_avatar(fn)}",
}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
assert res["data"]["avatar"] == f"data:image/png;base64,{encode_avatar(fn)}", res
@pytest.mark.p2
def test_description(self, WebApiAuth, add_dataset_func):
kb_id = add_dataset_func
payload = {"name": "description", "description": "description", "chunk_method": "naive"}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
assert res["data"]["description"] == "description", res
@pytest.mark.p1
@pytest.mark.parametrize(
"embedding_model",
[
"BAAI/bge-small-en-v1.5@Builtin",
"embedding-3@ZHIPU-AI",
],
ids=["builtin_baai", "tenant_zhipu"],
)
def test_embedding_model(self, WebApiAuth, add_dataset_func, embedding_model):
kb_id = add_dataset_func
payload = {"name": "embedding_model", "description": "", "chunk_method": "naive", "embedding_model": embedding_model}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
assert res["data"]["embedding_model"] == embedding_model, res
@pytest.mark.p2
@pytest.mark.parametrize(
"permission",
[
"me",
"team",
],
ids=["me", "team"],
)
def test_permission(self, WebApiAuth, add_dataset_func, permission):
kb_id = add_dataset_func
payload = {"name": "permission", "description": "", "chunk_method": "naive", "permission": permission}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
assert res["data"]["permission"] == permission.lower().strip(), res
@pytest.mark.p1
@pytest.mark.parametrize(
"chunk_method",
[
"naive",
"book",
"email",
"laws",
"manual",
"one",
"paper",
"picture",
"presentation",
"qa",
"table",
pytest.param("tag", marks=pytest.mark.skipif(os.getenv("DOC_ENGINE") == "infinity", reason="Infinity does not support parser_id=tag")),
],
ids=["naive", "book", "email", "laws", "manual", "one", "paper", "picture", "presentation", "qa", "table", "tag"],
)
def test_chunk_method(self, WebApiAuth, add_dataset_func, chunk_method):
kb_id = add_dataset_func
payload = {"name": "chunk_method", "description": "", "chunk_method": chunk_method}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
assert res["data"]["chunk_method"] == chunk_method, res
@pytest.mark.p1
@pytest.mark.skipif(os.getenv("DOC_ENGINE") != "infinity", reason="Infinity does not support parser_id=tag")
def test_chunk_method_tag_with_infinity(self, WebApiAuth, add_dataset_func):
kb_id = add_dataset_func
payload = {"name": "chunk_method", "description": "", "chunk_method": "tag"}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 103, res
assert res["message"] == "The chunking method Tag has not been supported by Infinity yet.", res
@pytest.mark.skipif(os.getenv("DOC_ENGINE") == "infinity", reason="#8208")
@pytest.mark.p2
@pytest.mark.parametrize("pagerank", [0, 50, 100], ids=["min", "mid", "max"])
def test_pagerank(self, WebApiAuth, add_dataset_func, pagerank):
kb_id = add_dataset_func
payload = {"name": "pagerank", "description": "", "chunk_method": "naive", "pagerank": pagerank}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
assert res["data"]["pagerank"] == pagerank, res
@pytest.mark.skipif(os.getenv("DOC_ENGINE") == "infinity", reason="#8208")
@pytest.mark.p2
def test_pagerank_set_to_0(self, WebApiAuth, add_dataset_func):
kb_id = add_dataset_func
payload = {"name": "pagerank", "description": "", "chunk_method": "naive", "pagerank": 50}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
assert res["data"]["pagerank"] == 50, res
payload = {"name": "pagerank", "description": "", "chunk_method": "naive", "pagerank": 0}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
assert res["data"]["pagerank"] == 0, res
@pytest.mark.skipif(os.getenv("DOC_ENGINE") != "infinity", reason="#8208")
@pytest.mark.p2
def test_pagerank_infinity(self, WebApiAuth, add_dataset_func):
kb_id = add_dataset_func
payload = {"name": "pagerank", "description": "", "chunk_method": "naive", "pagerank": 50}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 102, res
assert res["message"] == "'pagerank' can only be set when doc_engine is elasticsearch", res
@pytest.mark.p1
@pytest.mark.parametrize(
"parser_config",
[
{"auto_keywords": 0},
{"auto_keywords": 16},
{"auto_keywords": 32},
{"auto_questions": 0},
{"auto_questions": 5},
{"auto_questions": 10},
{"chunk_token_num": 1},
{"chunk_token_num": 1024},
{"chunk_token_num": 2048},
{"delimiter": "\n"},
{"delimiter": " "},
{"html4excel": True},
{"html4excel": False},
{"layout_recognize": "DeepDOC"},
{"layout_recognize": "Plain Text"},
{"tag_kb_ids": ["1", "2"]},
{"topn_tags": 1},
{"topn_tags": 5},
{"topn_tags": 10},
{"filename_embd_weight": 0.1},
{"filename_embd_weight": 0.5},
{"filename_embd_weight": 1.0},
{"task_page_size": 1},
{"task_page_size": None},
{"pages": [[1, 100]]},
{"pages": None},
{"graphrag": {"use_graphrag": True}},
{"graphrag": {"use_graphrag": False}},
{"graphrag": {"entity_types": ["age", "sex", "height", "weight"]}},
{"graphrag": {"method": "general"}},
{"graphrag": {"method": "light"}},
{"graphrag": {"community": True}},
{"graphrag": {"community": False}},
{"graphrag": {"resolution": True}},
{"graphrag": {"resolution": False}},
{"raptor": {"use_raptor": True}},
{"raptor": {"use_raptor": False}},
{"raptor": {"prompt": "Who are you?"}},
{"raptor": {"max_token": 1}},
{"raptor": {"max_token": 1024}},
{"raptor": {"max_token": 2048}},
{"raptor": {"threshold": 0.0}},
{"raptor": {"threshold": 0.5}},
{"raptor": {"threshold": 1.0}},
{"raptor": {"max_cluster": 1}},
{"raptor": {"max_cluster": 512}},
{"raptor": {"max_cluster": 1024}},
{"raptor": {"random_seed": 0}},
],
ids=[
"auto_keywords_min",
"auto_keywords_mid",
"auto_keywords_max",
"auto_questions_min",
"auto_questions_mid",
"auto_questions_max",
"chunk_token_num_min",
"chunk_token_num_mid",
"chunk_token_num_max",
"delimiter",
"delimiter_space",
"html4excel_true",
"html4excel_false",
"layout_recognize_DeepDOC",
"layout_recognize_navie",
"tag_kb_ids",
"topn_tags_min",
"topn_tags_mid",
"topn_tags_max",
"filename_embd_weight_min",
"filename_embd_weight_mid",
"filename_embd_weight_max",
"task_page_size_min",
"task_page_size_None",
"pages",
"pages_none",
"graphrag_true",
"graphrag_false",
"graphrag_entity_types",
"graphrag_method_general",
"graphrag_method_light",
"graphrag_community_true",
"graphrag_community_false",
"graphrag_resolution_true",
"graphrag_resolution_false",
"raptor_true",
"raptor_false",
"raptor_prompt",
"raptor_max_token_min",
"raptor_max_token_mid",
"raptor_max_token_max",
"raptor_threshold_min",
"raptor_threshold_mid",
"raptor_threshold_max",
"raptor_max_cluster_min",
"raptor_max_cluster_mid",
"raptor_max_cluster_max",
"raptor_random_seed_min",
],
)
def test_parser_config(self, WebApiAuth, add_dataset_func, parser_config):
kb_id = add_dataset_func
payload = {"name": "parser_config", "description": "", "chunk_method": "naive", "parser_config": parser_config}
res = update_dataset(WebApiAuth, kb_id, payload)
assert res["code"] == 0, res
for key, value in parser_config.items():
if not isinstance(value, dict):
assert res["data"]["parser_config"].get(key) == value, res
else:
for sub_key, sub_value in value.items():
assert res["data"]["parser_config"].get(key, {}).get(sub_key) == sub_value, res
@pytest.mark.p2
@pytest.mark.parametrize(
"payload",
[
{"id": "id"},
{"tenant_id": "e57c1966f99211efb41e9e45646e0111"},
{"created_by": "created_by"},
{"create_date": "Tue, 11 Mar 2025 13:37:23 GMT"},
{"create_time": 1741671443322},
{"update_date": "Tue, 11 Mar 2025 13:37:23 GMT"},
{"update_time": 1741671443339},
],
)
def test_field_unsupported(self, WebApiAuth, add_dataset_func, payload):
kb_id = add_dataset_func
full_payload = {"name": "field_unsupported", "description": "", "chunk_method": "naive", **payload}
res = update_dataset(WebApiAuth, kb_id, full_payload)
assert res["code"] == 101, res
assert "are not permitted" in res["message"], res