|
|
|
@@ -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 |