# # Copyright 2026 The InfiniFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import string import pytest from test_common import list_datasets, upload_documents from configs import DOCUMENT_NAME_LIMIT, INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth from utils.file_utils import create_txt_file from concurrent.futures import ThreadPoolExecutor, as_completed @pytest.mark.p1 @pytest.mark.usefixtures("clear_datasets") class TestAuthorization: @pytest.mark.parametrize( "invalid_auth, expected_code, expected_message", [ (None, 401, ""), (RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, ""), ], ) def test_invalid_auth(self, invalid_auth, expected_code, expected_message): res = upload_documents(invalid_auth) assert res["code"] == expected_code, res assert res["message"] == expected_message, res class TestDocumentsUpload: @pytest.mark.p1 def test_valid_single_upload(self, WebApiAuth, add_dataset_func, tmp_path): kb_id = add_dataset_func fp = create_txt_file(tmp_path / "ragflow_test.txt") res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp]) assert res["code"] == 0, res # New API returns "dataset_id" instead of "kb_id" due to key mapping assert res["data"][0]["dataset_id"] == kb_id, res assert res["data"][0]["name"] == fp.name, res @pytest.mark.p1 @pytest.mark.parametrize( "generate_test_files", [ "docx", "excel", "ppt", "image", "pdf", "txt", "md", "json", "eml", "html", ], indirect=True, ) def test_file_type_validation(self, WebApiAuth, add_dataset_func, generate_test_files, request): kb_id = add_dataset_func fp = generate_test_files[request.node.callspec.params["generate_test_files"]] res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp]) assert res["code"] == 0, res # New API returns "dataset_id" instead of "kb_id" due to key mapping assert res["data"][0]["dataset_id"] == kb_id, res assert res["data"][0]["name"] == fp.name, res @pytest.mark.p3 @pytest.mark.parametrize( "file_type", ["exe", "unknown"], ) def test_unsupported_file_type(self, WebApiAuth, add_dataset_func, tmp_path, file_type): kb_id = add_dataset_func fp = tmp_path / f"ragflow_test.{file_type}" fp.touch() res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp]) assert res["code"] == 500, res assert res["message"] == f"ragflow_test.{file_type}: This type of file has not been supported yet!", res @pytest.mark.p2 def test_missing_file(self, WebApiAuth, add_dataset_func): kb_id = add_dataset_func res = upload_documents(WebApiAuth, {"kb_id": kb_id}) assert res["code"] == 101, res assert res["message"] == "No file part!", res @pytest.mark.p3 def test_empty_file(self, WebApiAuth, add_dataset_func, tmp_path): kb_id = add_dataset_func fp = tmp_path / "empty.txt" fp.touch() res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp]) assert res["code"] == 0, res assert res["data"][0]["size"] == 0, res @pytest.mark.p3 def test_filename_empty(self, WebApiAuth, add_dataset_func, tmp_path): kb_id = add_dataset_func fp = create_txt_file(tmp_path / "ragflow_test.txt") res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp], filename_override="") assert res["code"] == 101, res assert res["message"] == "No file selected!", res @pytest.mark.p3 def test_filename_exceeds_max_length(self, WebApiAuth, add_dataset_func, tmp_path): kb_id = add_dataset_func fp = create_txt_file(tmp_path / f"{'a' * (DOCUMENT_NAME_LIMIT - 4)}.txt") res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp]) assert res["code"] == 0, res assert res["data"][0]["name"] == fp.name, res @pytest.mark.p2 def test_invalid_kb_id(self, WebApiAuth, tmp_path): fp = create_txt_file(tmp_path / "ragflow_test.txt") res = upload_documents(WebApiAuth, {"kb_id": "invalid_kb_id"}, [fp]) assert res["code"] == 102, res assert res["message"] == "Can't find the dataset with ID invalid_kb_id!", res @pytest.mark.p2 def test_duplicate_files(self, WebApiAuth, add_dataset_func, tmp_path): kb_id = add_dataset_func fp = create_txt_file(tmp_path / "ragflow_test.txt") res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp, fp]) assert res["code"] == 0, res assert len(res["data"]) == 2, res for i in range(len(res["data"])): # New API returns "dataset_id" instead of "kb_id" due to key mapping assert res["data"][i]["dataset_id"] == kb_id, res expected_name = fp.name if i != 0: expected_name = f"{fp.stem}({i}){fp.suffix}" assert res["data"][i]["name"] == expected_name, res @pytest.mark.p3 def test_filename_special_characters(self, WebApiAuth, add_dataset_func, tmp_path): kb_id = add_dataset_func illegal_chars = '<>:"/\\|?*' translation_table = str.maketrans({char: "_" for char in illegal_chars}) safe_filename = string.punctuation.translate(translation_table) fp = tmp_path / f"{safe_filename}.txt" fp.write_text("Sample text content") res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp]) assert res["code"] == 0, res assert len(res["data"]) == 1, res # New API returns "dataset_id" instead of "kb_id" due to key mapping assert res["data"][0]["dataset_id"] == kb_id, res assert res["data"][0]["name"] == fp.name, res @pytest.mark.p1 def test_multiple_files(self, WebApiAuth, add_dataset_func, tmp_path): kb_id = add_dataset_func expected_document_count = 20 fps = [] for i in range(expected_document_count): fp = create_txt_file(tmp_path / f"ragflow_test_{i}.txt") fps.append(fp) res = upload_documents(WebApiAuth, {"kb_id": kb_id}, fps) assert res["code"] == 0, res res = list_datasets(WebApiAuth) assert res["data"][0]["document_count"] == expected_document_count, res @pytest.mark.p3 def test_concurrent_upload(self, WebApiAuth, add_dataset_func, tmp_path): kb_id = add_dataset_func count = 20 fps = [] for i in range(count): fp = create_txt_file(tmp_path / f"ragflow_test_{i}.txt") fps.append(fp) with ThreadPoolExecutor(max_workers=5) as executor: futures = [executor.submit(upload_documents, WebApiAuth, {"kb_id": kb_id}, fps[i : i + 1]) 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), responses res = list_datasets(WebApiAuth) assert res["data"][0]["document_count"] == count, res import asyncio from types import SimpleNamespace class _AwaitableValue: def __init__(self, value): self._value = value def __await__(self): async def _coro(): return self._value return _coro().__await__() class _DummyFiles(dict): def getlist(self, key): value = self.get(key, []) if isinstance(value, list): return value return [value] class _DummyFile: def __init__(self, filename): self.filename = filename self.closed = False self.stream = self def close(self): self.closed = True class _DummyRequest: def __init__(self, form=None, files=None, args=None): self._form = form or {} self._files = files or _DummyFiles() self.args = args or {} @property def form(self): return _AwaitableValue(self._form) @property def files(self): return _AwaitableValue(self._files) def _run(coro): return asyncio.run(coro) @pytest.mark.p2 class TestDocumentsUploadUnit: """Unit tests for document upload using upload_documents helper function""" def test_missing_kb_id(self, WebApiAuth, tmp_path): """Test that missing KB ID returns error""" # When kb_id is empty, the API should return an error fp = create_txt_file(tmp_path / "ragflow_test.txt") res = upload_documents(WebApiAuth, {"kb_id": ""}, [fp]) assert res["code"] == 100 assert res["message"] == "" def test_missing_file_part(self, WebApiAuth, add_dataset_func): """Test that missing file part returns error""" kb_id = add_dataset_func # Call without files - should return error for missing file res = upload_documents(WebApiAuth, {"kb_id": kb_id}) assert res["code"] == 101 assert "file" in res["message"].lower() def test_empty_filename_closes_files(self, WebApiAuth, add_dataset_func, tmp_path): """Test that empty filename returns error""" kb_id = add_dataset_func # Create a file with empty name by using filename_override fp = create_txt_file(tmp_path / "ragflow_test.txt") res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp], filename_override="") assert res["code"] == 101 assert "file" in res["message"].lower() or "selected" in res["message"].lower() def test_invalid_kb_id_raises(self, WebApiAuth, tmp_path): """Test that invalid KB ID returns error""" fp = create_txt_file(tmp_path / "ragflow_test.txt") res = upload_documents(WebApiAuth, {"kb_id": "invalid_kb_id"}, [fp]) # The API should return an error for invalid KB ID assert res["code"] == 102 assert "Can't find the dataset" in res["message"] or "not found" in res["message"].lower() def test_no_permission(self, WebApiAuth, tmp_path): """Test that no permission returns error""" # Create a file and try to upload to a dataset we don't have access to # This test would require setting up a dataset without permission # For now, we skip this test as it requires specific setup pytest.skip("Requires dataset without permission setup") def test_thread_pool_errors(self, WebApiAuth, add_dataset_func, tmp_path): """Test that thread pool errors are handled""" kb_id = add_dataset_func # Upload a file with unsupported type fp = tmp_path / "test.exe" fp.write_text("test") res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp]) # Should return error for unsupported file type assert res["code"] == 500 assert "supported" in res["message"].lower() or "type" in res["message"].lower() def test_empty_upload_result(self, WebApiAuth, add_dataset_func, tmp_path): """Test that empty upload result returns error""" kb_id = add_dataset_func # Create an empty file fp = tmp_path / "empty.txt" fp.write_text("") res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp]) # Empty file might cause issues # The exact behavior depends on the implementation # Just verify we get a response assert "code" in res @pytest.mark.p2 class TestWebCrawlUnit: def test_invalid_url(self, document_rest_api_module, monkeypatch): module = document_rest_api_module monkeypatch.setattr( module, "request", _DummyRequest(form={"name": "doc", "url": "not-a-url"}, args={"type": "web"}), ) res = _run(module.upload_document(dataset_id="kb1")) assert res["code"] == 101 assert res["message"] == "The URL format is invalid" def test_invalid_kb_id(self, document_rest_api_module, monkeypatch): module = document_rest_api_module monkeypatch.setattr(module, "is_valid_url", lambda _url: True) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None)) monkeypatch.setattr( module, "request", _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}), ) res = _run(module.upload_document(dataset_id="missing")) assert res["code"] == 102 assert "Can't find the dataset" in res["message"] def test_no_permission(self, document_rest_api_module, monkeypatch): module = document_rest_api_module kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) monkeypatch.setattr(module, "is_valid_url", lambda _url: True) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: False) monkeypatch.setattr( module, "request", _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}), ) res = _run(module.upload_document(dataset_id="kb1")) assert res["code"] == 109 assert res["message"] == "No authorization." def test_download_failure(self, document_rest_api_module, monkeypatch): module = document_rest_api_module kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) monkeypatch.setattr(module, "is_valid_url", lambda _url: True) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) monkeypatch.setattr(module, "html2pdf", lambda _url: None) monkeypatch.setattr( module, "request", _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}), ) res = _run(module.upload_document(dataset_id="kb1")) assert res["code"] == 100 assert "Download failure" in res["message"] def test_unsupported_type(self, document_rest_api_module, monkeypatch): module = document_rest_api_module kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) monkeypatch.setattr(module, "is_valid_url", lambda _url: True) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4") monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"}) monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None) monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"}) monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"}) monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: "bad.exe") monkeypatch.setattr( module, "request", _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}), ) res = _run(module.upload_document(dataset_id="kb1")) assert res["code"] == 100 assert "supported yet" in res["message"] @pytest.mark.parametrize( "filename,filetype,expected_parser", [ ("image.png", "visual", "picture"), ("sound.mp3", "aural", "audio"), ("deck.pptx", "doc", "presentation"), ("mail.eml", "doc", "email"), ], ) def test_success_parser_overrides(self, document_rest_api_module, monkeypatch, filename, filetype, expected_parser): module = document_rest_api_module kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) captured = {} class _Storage: def obj_exist(self, *_args, **_kwargs): return False def put(self, *_args, **_kwargs): captured["put"] = True def insert_doc(doc): captured["doc"] = doc monkeypatch.setattr(module, "is_valid_url", lambda _url: True) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4") monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"}) monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None) monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"}) monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"}) monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: filename) monkeypatch.setattr(module, "filename_type", lambda _name: filetype) monkeypatch.setattr(module, "thumbnail", lambda *_args, **_kwargs: "") monkeypatch.setattr(module, "get_uuid", lambda: "doc-1") monkeypatch.setattr(module.settings, "STORAGE_IMPL", _Storage()) monkeypatch.setattr(module.DocumentService, "insert", insert_doc) monkeypatch.setattr(module.FileService, "add_file_from_kb", lambda *_args, **_kwargs: None) monkeypatch.setattr( module, "request", _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}), ) res = _run(module.upload_document(dataset_id="kb1")) assert res["code"] == 0 assert captured["doc"]["parser_id"] == expected_parser assert captured["put"] is True def test_exception_path(self, document_rest_api_module, monkeypatch): module = document_rest_api_module kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={}) class _Storage: def obj_exist(self, *_args, **_kwargs): return False def put(self, *_args, **_kwargs): return None def insert_doc(_doc): raise RuntimeError("boom") monkeypatch.setattr(module, "is_valid_url", lambda _url: True) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True) monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4") monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"}) monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None) monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"}) monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"}) monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: "doc.pdf") monkeypatch.setattr(module, "filename_type", lambda _name: "pdf") monkeypatch.setattr(module, "thumbnail", lambda *_args, **_kwargs: "") monkeypatch.setattr(module, "get_uuid", lambda: "doc-1") monkeypatch.setattr(module.settings, "STORAGE_IMPL", _Storage()) monkeypatch.setattr(module.DocumentService, "insert", insert_doc) monkeypatch.setattr(module.FileService, "add_file_from_kb", lambda *_args, **_kwargs: None) monkeypatch.setattr( module, "request", _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}), ) res = _run(module.upload_document(dataset_id="kb1")) assert res["code"] == 100