diff --git a/sdk/python/test/test_http_api/common.py b/sdk/python/test/test_http_api/common.py index d8d7b1394..5bcb77ffb 100644 --- a/sdk/python/test/test_http_api/common.py +++ b/sdk/python/test/test_http_api/common.py @@ -191,8 +191,8 @@ def list_chunks(auth, dataset_id, document_id, params=None): def batch_add_chunks(auth, dataset_id, document_id, num): - ids = [] + chunk_ids = [] for i in range(num): res = add_chunk(auth, dataset_id, document_id, {"content": f"ragflow test {i}"}) - ids.append(res["data"]["chunk"]["id"]) - return ids + chunk_ids.append(res["data"]["chunk"]["id"]) + return chunk_ids diff --git a/sdk/python/test/test_http_api/test_chunk_management_within_dataset/test_add_chunk.py b/sdk/python/test/test_http_api/test_chunk_management_within_dataset/test_add_chunk.py index 683f84571..338a0ade1 100644 --- a/sdk/python/test/test_http_api/test_chunk_management_within_dataset/test_add_chunk.py +++ b/sdk/python/test/test_http_api/test_chunk_management_within_dataset/test_add_chunk.py @@ -16,7 +16,7 @@ from concurrent.futures import ThreadPoolExecutor import pytest -from common import INVALID_API_TOKEN, add_chunk, delete_documnet +from common import INVALID_API_TOKEN, add_chunk, delete_documnet, list_chunks from libs.auth import RAGFlowHttpApiAuth @@ -69,10 +69,14 @@ class TestAddChunk: ) def test_content(self, get_http_api_auth, get_dataset_id_and_document_id, payload, expected_code, expected_message): dataset_id, document_id = get_dataset_id_and_document_id + res = list_chunks(get_http_api_auth, dataset_id, document_id) + chunks_count = res["data"]["doc"]["chunk_count"] res = add_chunk(get_http_api_auth, dataset_id, document_id, payload) assert res["code"] == expected_code if expected_code == 0: validate_chunk_details(dataset_id, document_id, payload, res) + res = list_chunks(get_http_api_auth, dataset_id, document_id) + assert res["data"]["doc"]["chunk_count"] == chunks_count + 1 else: assert res["message"] == expected_message @@ -93,10 +97,14 @@ class TestAddChunk: ) def test_important_keywords(self, get_http_api_auth, get_dataset_id_and_document_id, payload, expected_code, expected_message): dataset_id, document_id = get_dataset_id_and_document_id + res = list_chunks(get_http_api_auth, dataset_id, document_id) + chunks_count = res["data"]["doc"]["chunk_count"] res = add_chunk(get_http_api_auth, dataset_id, document_id, payload) assert res["code"] == expected_code if expected_code == 0: validate_chunk_details(dataset_id, document_id, payload, res) + res = list_chunks(get_http_api_auth, dataset_id, document_id) + assert res["data"]["doc"]["chunk_count"] == chunks_count + 1 else: assert res["message"] == expected_message @@ -118,10 +126,14 @@ class TestAddChunk: ) def test_questions(self, get_http_api_auth, get_dataset_id_and_document_id, payload, expected_code, expected_message): dataset_id, document_id = get_dataset_id_and_document_id + res = list_chunks(get_http_api_auth, dataset_id, document_id) + chunks_count = res["data"]["doc"]["chunk_count"] res = add_chunk(get_http_api_auth, dataset_id, document_id, payload) assert res["code"] == expected_code if expected_code == 0: validate_chunk_details(dataset_id, document_id, payload, res) + res = list_chunks(get_http_api_auth, dataset_id, document_id) + assert res["data"]["doc"]["chunk_count"] == chunks_count + 1 else: assert res["message"] == expected_message @@ -169,13 +181,19 @@ class TestAddChunk: def test_repeated_add_chunk(self, get_http_api_auth, get_dataset_id_and_document_id): payload = {"content": "a"} dataset_id, document_id = get_dataset_id_and_document_id + res = list_chunks(get_http_api_auth, dataset_id, document_id) + chunks_count = res["data"]["doc"]["chunk_count"] res = add_chunk(get_http_api_auth, dataset_id, document_id, payload) assert res["code"] == 0 validate_chunk_details(dataset_id, document_id, payload, res) + res = list_chunks(get_http_api_auth, dataset_id, document_id) + assert res["data"]["doc"]["chunk_count"] == chunks_count + 1 res = add_chunk(get_http_api_auth, dataset_id, document_id, payload) assert res["code"] == 0 validate_chunk_details(dataset_id, document_id, payload, res) + res = list_chunks(get_http_api_auth, dataset_id, document_id) + assert res["data"]["doc"]["chunk_count"] == chunks_count + 2 def test_add_chunk_to_deleted_document(self, get_http_api_auth, get_dataset_id_and_document_id): dataset_id, document_id = get_dataset_id_and_document_id @@ -188,6 +206,8 @@ class TestAddChunk: def test_concurrent_add_chunk(self, get_http_api_auth, get_dataset_id_and_document_id): chunk_num = 50 dataset_id, document_id = get_dataset_id_and_document_id + res = list_chunks(get_http_api_auth, dataset_id, document_id) + chunks_count = res["data"]["doc"]["chunk_count"] with ThreadPoolExecutor(max_workers=5) as executor: futures = [ @@ -202,3 +222,5 @@ class TestAddChunk: ] responses = [f.result() for f in futures] assert all(r["code"] == 0 for r in responses) + res = list_chunks(get_http_api_auth, dataset_id, document_id) + assert res["data"]["doc"]["chunk_count"] == chunks_count + chunk_num