# # Copyright 2026 The InfiniFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import pytest from test.testcases.configs import IS_GO_PROXY from libs.auth import RAGFlowHttpApiAuth from test.testcases.restful_api.helpers.client import RestClient from utils.file_utils import create_txt_file from utils import wait_for GO_ONLY_SKIPS = { "Go route is not implemented": { "test_document_download_by_id_invalid_id_contract", "test_llm_factories_live_auth_contract", "test_llm_list_live_auth_contract", "test_retrieval_compatibility_endpoint", "test_retrieval_compatibility_requires_dataset_ids", "test_retrieval_compatibility_requires_auth", "test_retrieval_requires_auth_contract", "test_system_oceanbase_status_auth_contract", "test_task_routes_require_auth", "test_patch_task_rejects_unsupported_action", "test_cancel_missing_task_sets_cancel_contract", }, "Go validation or response contract does not match the established API contract": { "test_chat_create_name_validation", "test_chat_duplicate_name_validation", "test_chat_create_additional_guards_contract", "test_chat_list_default_get_and_separate_lookup_contract", "test_chat_list_page_and_page_size_contract", "test_chat_list_sorting_contract", "test_chat_create_prompt_contract", "test_chat_update_name_contract", "test_chat_update_mapping_and_validation_branches_p2", "test_dataset_update_name_and_case_insensitive_contract", "test_dataset_update_embedding_model_format_contract", "test_dataset_update_parser_config_valid_matrix_contract", "test_dataset_update_parser_config_with_chunk_method_change_contract", "test_dataset_update_pagerank_contract", "test_dataset_update_pagerank_set_to_zero_contract", "test_dataset_update_content_type_and_payload_contract", "test_dataset_update_identifier_validation_contract", "test_dataset_update_parser_config_defaults_contract", "test_dataset_update_parser_config_invalid_contract", "test_dataset_update_field_unset_and_unsupported_contract", "test_dataset_create_parser_config_missing_raptor_and_graphrag", "test_dataset_create_embedding_model_contract", "test_dataset_create_embedding_model_format_contract", "test_dataset_create_parser_config_bugfix_contract", "test_dataset_create_content_type_and_payload_bad_contract", "test_dataset_create_parser_config_invalid_contract", "test_dataset_create_parser_config_defaults_and_extra_fields_contract", "test_dataset_list_query_contract_matrix", "test_dataset_delete_contract_matrix", "test_memory_update_invalid_name", "test_memory_crud_cycle", "test_messages_add_list_recent_content_update_forget", "test_message_status_validation_requires_boolean", "test_message_search_route_contract", "test_memory_crud_and_config", "test_messages_list_and_search_validation_contracts", "test_message_update_forget_and_content_error_contracts", "test_search_create_invalid_name", "test_search_update_invalid_search_id", "test_session_create_validation_and_deleted_chat_contract", "test_session_delete_basic_scenarios", "test_session_list_filter_and_deleted_chat_contract", "test_session_list_page_and_sort_contract", "test_session_update_name_and_param_contract", "test_session_update_requires_auth_and_invalid_target_contract", "test_chat_completion_validation_errors", "test_chat_completion_nonstream_with_session", "test_chat_completion_nonstream_with_chat_without_session", "test_chat_completion_nonstream_without_chat", "test_chat_completion_stream_events", "test_search_completion_sse_shape_when_kb_ids_provided", "test_system_tokens_auth_and_crud", }, "Go LLM setup cannot exercise the configured model": { "test_related_questions_contract", }, } def pytest_collection_modifyitems(items): if not IS_GO_PROXY: return for item in items: test_name = item.name.split("[", 1)[0] for reason, skipped_tests in GO_ONLY_SKIPS.items(): if test_name in skipped_tests: item.add_marker(pytest.mark.skip(reason=reason)) break @pytest.fixture(scope="session") def RestApiAuth(token): return RAGFlowHttpApiAuth(token) @pytest.fixture(scope="session") def rest_client(token): return RestClient(token=token) @pytest.fixture(scope="session") def rest_client_noauth(): return RestClient(token=None) @pytest.fixture def clear_datasets(rest_client): def _cleanup(): res = rest_client.delete("/datasets", json={"ids": None, "delete_all": True}) assert res.status_code == 200, res.text payload = res.json() assert payload["code"] in (0, 102), payload yield _cleanup() @pytest.fixture def clear_chats(rest_client): def _cleanup(): res = rest_client.delete("/chats", json={"ids": None, "delete_all": True}) assert res.status_code == 200, res.text payload = res.json() assert payload["code"] in (0, 102), payload yield _cleanup() @pytest.fixture def create_dataset(rest_client, clear_datasets): created_ids: list[str] = [] def _create(name: str = "restful_dataset") -> str: res = rest_client.post("/datasets", json={"name": name}) assert res.status_code == 200 payload = res.json() assert payload["code"] == 0, payload dataset_id = payload["data"]["id"] created_ids.append(dataset_id) return dataset_id yield _create if created_ids: res = rest_client.delete("/datasets", json={"ids": created_ids}) assert res.status_code == 200 payload = res.json() # Dataset may already be removed by test logic/cleanup. assert payload["code"] in (0, 102), payload @pytest.fixture def create_chat(rest_client, clear_chats): created_ids: list[str] = [] def _create(name: str = "restful_chat") -> str: res = rest_client.post("/chats", json={"name": name, "dataset_ids": []}) assert res.status_code == 200 payload = res.json() assert payload["code"] == 0, payload chat_id = payload["data"]["id"] created_ids.append(chat_id) return chat_id yield _create if created_ids: res = rest_client.delete("/chats", json={"ids": created_ids}) assert res.status_code == 200, res.text payload = res.json() assert payload["code"] in (0, 102), payload @pytest.fixture def create_document(rest_client, create_dataset, tmp_path): created_docs: list[tuple[str, str]] = [] def _create(name: str = "restful_doc.txt") -> tuple[str, str]: dataset_id = create_dataset("dataset_for_doc") fp = create_txt_file(tmp_path / name) with fp.open("rb") as file_obj: files = [("file", (fp.name, file_obj))] res = rest_client.post(f"/datasets/{dataset_id}/documents", files=files) assert res.status_code == 200 payload = res.json() assert payload["code"] == 0, payload document_id = payload["data"][0]["id"] created_docs.append((dataset_id, document_id)) return dataset_id, document_id yield _create for dataset_id, document_id in created_docs: res = rest_client.delete(f"/datasets/{dataset_id}/documents", json={"ids": [document_id]}) assert res.status_code == 200, res.text payload = res.json() assert payload["code"] in (0, 102), payload @wait_for(60, 1, "Document parsing timeout in RESTful batch2 tests") def _parsed(rest_client: RestClient, dataset_id: str, document_id: str): res = rest_client.get(f"/datasets/{dataset_id}/documents", params={"id": document_id}) if res.status_code != 200: return False payload = res.json() if payload["code"] != 0: return False docs = payload["data"]["docs"] if not docs: return False return docs[0].get("run") == "DONE" @pytest.fixture def ensure_parsed_document(rest_client, create_document): def _ensure() -> tuple[str, str]: dataset_id, document_id = create_document() res = rest_client.post( f"/datasets/{dataset_id}/documents/parse", json={"document_ids": [document_id]}, ) assert res.status_code == 200 payload = res.json() assert payload["code"] == 0, payload _parsed(rest_client, dataset_id, document_id) return dataset_id, document_id return _ensure @pytest.hookimpl(tryfirst=True, hookwrapper=True) def pytest_runtest_protocol(item, nextitem): import time start = time.perf_counter() yield duration = time.perf_counter() - start if duration >= 10: color = "****" # 4 stars elif duration >= 5: color = "***" # 3 stars elif duration >= 1: color = "**" # 2 stars else: color = "" print(f" {color} [{duration:.3f}s]")