File size: 3,708 Bytes
4359d15 090b2e7 4359d15 090b2e7 4359d15 090b2e7 4359d15 090b2e7 4359d15 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
from test_sdkbase import TestSdk
from ragflow import RAGFlow
import pytest
from common import API_KEY, HOST_ADDRESS
class TestDataset(TestSdk):
def test_create_dataset(self):
'''
1. create a kb
2. list the kb
3. get the detail info according to the kb id
4. update the kb
5. delete the kb
'''
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
# create a kb
res = ragflow.create_dataset("kb1")
assert res['code'] == 0 and res['message'] == 'success'
dataset_name = res['data']['dataset_name']
def test_list_dataset_success(self):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
# Call the list_datasets method
response = ragflow.list_dataset()
code, datasets = response
assert code == 200
def test_list_dataset_with_checking_size_and_name(self):
datasets_to_create = ["dataset1", "dataset2", "dataset3"]
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
created_response = [ragflow.create_dataset(name) for name in datasets_to_create]
real_name_to_create = set()
for response in created_response:
assert 'data' in response, "Response is missing 'data' key"
dataset_name = response['data']['dataset_name']
real_name_to_create.add(dataset_name)
status_code, listed_data = ragflow.list_dataset(0, 3)
listed_data = listed_data['data']
listed_names = {d['name'] for d in listed_data}
assert listed_names == real_name_to_create
assert status_code == 200
assert len(listed_data) == len(datasets_to_create)
def test_list_dataset_with_getting_empty_result(self):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
datasets_to_create = []
created_response = [ragflow.create_dataset(name) for name in datasets_to_create]
real_name_to_create = set()
for response in created_response:
assert 'data' in response, "Response is missing 'data' key"
dataset_name = response['data']['dataset_name']
real_name_to_create.add(dataset_name)
status_code, listed_data = ragflow.list_dataset(0, 0)
listed_data = listed_data['data']
listed_names = {d['name'] for d in listed_data}
assert listed_names == real_name_to_create
assert status_code == 200
assert len(listed_data) == 0
def test_list_dataset_with_creating_100_knowledge_bases(self):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
datasets_to_create = ["dataset1"] * 100
created_response = [ragflow.create_dataset(name) for name in datasets_to_create]
real_name_to_create = set()
for response in created_response:
assert 'data' in response, "Response is missing 'data' key"
dataset_name = response['data']['dataset_name']
real_name_to_create.add(dataset_name)
status_code, listed_data = ragflow.list_dataset(0, 100)
listed_data = listed_data['data']
listed_names = {d['name'] for d in listed_data}
assert listed_names == real_name_to_create
assert status_code == 200
assert len(listed_data) == 100
def test_list_dataset_with_showing_one_dataset(self):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
response = ragflow.list_dataset(0, 1)
code, response = response
datasets = response['data']
assert len(datasets) == 1
def test_list_dataset_failure(self):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
response = ragflow.list_dataset(-1, -1)
_, res = response
assert "IndexError" in res['message']
|