mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-04 04:49:08 +08:00
a963afa0d3
Signed-off-by: zhuwenxing <wenxing.zhu@zilliz.com>
1136 lines
43 KiB
Python
1136 lines
43 KiB
Python
import datetime
|
|
import logging
|
|
import time
|
|
from utils.util_log import test_log as logger
|
|
from utils.utils import gen_collection_name
|
|
import pytest
|
|
from api.milvus import CollectionClient
|
|
from base.testbase import TestBase
|
|
import threading
|
|
from utils.utils import get_data_by_payload
|
|
from pymilvus import (
|
|
FieldSchema, CollectionSchema, DataType,
|
|
Collection
|
|
)
|
|
|
|
|
|
@pytest.mark.L0
|
|
class TestCreateCollection(TestBase):
|
|
|
|
@pytest.mark.parametrize("dim", [128])
|
|
def test_create_collections_quick_setup(self, dim):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
}
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
assert rsp['data']['autoId'] is False
|
|
assert rsp['data']['enableDynamicField'] is True
|
|
assert "COSINE" in str(rsp['data']["indexes"])
|
|
|
|
@pytest.mark.parametrize("dim", [128])
|
|
@pytest.mark.parametrize("metric_type", ["L2", "COSINE", "IP"])
|
|
@pytest.mark.parametrize("id_type", ["Int64", "VarChar"])
|
|
@pytest.mark.parametrize("primary_field", ["id", "url"])
|
|
@pytest.mark.parametrize("vector_field", ["vector", "embedding"])
|
|
def test_create_collection_quick_setup_with_custom(self, vector_field, primary_field, dim, id_type, metric_type):
|
|
"""
|
|
Insert a vector with a simple payload
|
|
"""
|
|
# create a collection
|
|
name = gen_collection_name()
|
|
collection_payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
"metricType": metric_type,
|
|
"primaryFieldName": primary_field,
|
|
"vectorFieldName": vector_field,
|
|
"idType": id_type,
|
|
}
|
|
if id_type == "VarChar":
|
|
collection_payload["params"] = {"max_length": "256"}
|
|
rsp = self.collection_client.collection_create(collection_payload)
|
|
assert rsp['code'] == 0
|
|
rsp = self.collection_client.collection_describe(name)
|
|
logger.info(f"rsp: {rsp}")
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
fields = [f["name"] for f in rsp['data']['fields']]
|
|
assert primary_field in fields
|
|
assert vector_field in fields
|
|
for f in rsp['data']['fields']:
|
|
if f['name'] == primary_field:
|
|
assert f['type'] == id_type
|
|
assert f['primaryKey'] is True
|
|
for index in rsp['data']['indexes']:
|
|
assert index['metricType'] == metric_type
|
|
|
|
@pytest.mark.parametrize("enable_dynamic_field", [False, "False", "0"])
|
|
@pytest.mark.parametrize("request_shards_num", [2, "2"])
|
|
@pytest.mark.parametrize("request_ttl_seconds", [360, "360"])
|
|
def test_create_collections_without_params(self, enable_dynamic_field, request_shards_num, request_ttl_seconds):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
metric_type = "COSINE"
|
|
client = self.collection_client
|
|
num_shards = 2
|
|
consistency_level = "Strong"
|
|
ttl_seconds = 360
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
"metricType": metric_type,
|
|
"params":{
|
|
"enableDynamicField": enable_dynamic_field,
|
|
"shardsNum": request_shards_num,
|
|
"consistencyLevel": f"{consistency_level}",
|
|
"ttlSeconds": request_ttl_seconds,
|
|
},
|
|
}
|
|
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection by pymilvus
|
|
c = Collection(name)
|
|
res = c.describe()
|
|
logger.info(f"describe collection: {res}")
|
|
# describe collection
|
|
time.sleep(10)
|
|
rsp = client.collection_describe(name)
|
|
logger.info(f"describe collection: {rsp}")
|
|
|
|
ttl_seconds_actual = None
|
|
for d in rsp["data"]["properties"]:
|
|
if d["key"] == "collection.ttl.seconds":
|
|
ttl_seconds_actual = int(d["value"])
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['enableDynamicField'] == False
|
|
assert rsp['data']['collectionName'] == name
|
|
assert rsp['data']['shardsNum'] == num_shards
|
|
assert rsp['data']['consistencyLevel'] == consistency_level
|
|
assert ttl_seconds_actual == ttl_seconds
|
|
|
|
def test_create_collections_with_all_params(self):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
metric_type = "COSINE"
|
|
client = self.collection_client
|
|
num_shards = 2
|
|
num_partitions = 36
|
|
consistency_level = "Strong"
|
|
ttl_seconds = 360
|
|
payload = {
|
|
"collectionName": name,
|
|
"enableDynamicField": True,
|
|
"params":{
|
|
"shardsNum": f"{num_shards}",
|
|
"partitionsNum": f"{num_partitions}",
|
|
"consistencyLevel": f"{consistency_level}",
|
|
"ttlSeconds": f"{ttl_seconds}",
|
|
},
|
|
"schema": {
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "word_count", "dataType": "Int64", "isPartitionKey": True, "elementTypeParams": {}},
|
|
{"fieldName": "book_describe", "dataType": "VarChar", "elementTypeParams": {"max_length": "256"}},
|
|
{"fieldName": "json", "dataType": "JSON", "elementTypeParams": {}},
|
|
{"fieldName": "int_array", "dataType": "Array", "elementDataType": "Int64",
|
|
"elementTypeParams": {"max_capacity": "1024"}},
|
|
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}}
|
|
]
|
|
},
|
|
"indexParams": [
|
|
{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": f"{metric_type}"}]
|
|
}
|
|
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection by pymilvus
|
|
c = Collection(name)
|
|
res = c.describe()
|
|
logger.info(f"describe collection: {res}")
|
|
# describe collection
|
|
time.sleep(10)
|
|
rsp = client.collection_describe(name)
|
|
logger.info(f"describe collection: {rsp}")
|
|
|
|
ttl_seconds_actual = None
|
|
for d in rsp["data"]["properties"]:
|
|
if d["key"] == "collection.ttl.seconds":
|
|
ttl_seconds_actual = int(d["value"])
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
assert rsp['data']['shardsNum'] == num_shards
|
|
assert rsp['data']['partitionsNum'] == num_partitions
|
|
assert rsp['data']['consistencyLevel'] == consistency_level
|
|
assert ttl_seconds_actual == ttl_seconds
|
|
|
|
|
|
@pytest.mark.parametrize("auto_id", [True, False])
|
|
@pytest.mark.parametrize("enable_dynamic_field", [True, False])
|
|
@pytest.mark.parametrize("enable_partition_key", [True, False])
|
|
@pytest.mark.parametrize("dim", [128])
|
|
def test_create_collections_custom_without_index(self, dim, auto_id, enable_dynamic_field, enable_partition_key):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"schema": {
|
|
"autoId": auto_id,
|
|
"enableDynamicField": enable_dynamic_field,
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "user_id", "dataType": "Int64", "isPartitionKey": enable_partition_key,
|
|
"elementTypeParams": {}},
|
|
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
|
|
{"fieldName": "book_describe", "dataType": "VarChar", "elementTypeParams": {"max_length": "256"}},
|
|
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}},
|
|
{"fieldName": "image_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}},
|
|
]
|
|
}
|
|
}
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
c = Collection(name)
|
|
logger.info(f"schema: {c.schema}")
|
|
# describe collection
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
assert rsp['data']['autoId'] == auto_id
|
|
assert c.schema.auto_id == auto_id
|
|
assert rsp['data']['enableDynamicField'] == enable_dynamic_field
|
|
assert c.schema.enable_dynamic_field == enable_dynamic_field
|
|
# assert no index created
|
|
indexes = rsp['data']['indexes']
|
|
assert len(indexes) == 0
|
|
# assert not loaded
|
|
assert rsp['data']['load'] == "LoadStateNotLoad"
|
|
for field in rsp['data']['fields']:
|
|
if field['name'] == "user_id":
|
|
assert field['partitionKey'] == enable_partition_key
|
|
for field in c.schema.fields:
|
|
if field.name == "user_id":
|
|
assert field.is_partition_key == enable_partition_key
|
|
|
|
@pytest.mark.parametrize("metric_type", ["L2", "IP", "COSINE"])
|
|
@pytest.mark.parametrize("dim", [128])
|
|
def test_create_collections_one_float_vector_with_index(self, dim, metric_type):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"schema": {
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
|
|
{"fieldName": "book_describe", "dataType": "VarChar", "elementTypeParams": {"max_length": "256"}},
|
|
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}}
|
|
]
|
|
},
|
|
"indexParams": [
|
|
{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": f"{metric_type}"}]
|
|
}
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
time.sleep(10)
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
# assert index created
|
|
indexes = rsp['data']['indexes']
|
|
assert len(indexes) == len(payload['indexParams'])
|
|
# assert load success
|
|
assert rsp['data']['load'] == "LoadStateLoaded"
|
|
|
|
@pytest.mark.parametrize("metric_type", ["L2", "IP", "COSINE"])
|
|
@pytest.mark.parametrize("dim", [128])
|
|
def test_create_collections_multi_float_vector_with_one_index(self, dim, metric_type):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"schema": {
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
|
|
{"fieldName": "book_describe", "dataType": "VarChar", "elementTypeParams": {"max_length": "256"}},
|
|
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}},
|
|
{"fieldName": "image_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}}
|
|
]
|
|
},
|
|
"indexParams": [
|
|
{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": f"{metric_type}"}]
|
|
}
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 65535
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
time.sleep(10)
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
# assert index created
|
|
indexes = rsp['data']['indexes']
|
|
assert len(indexes) == len(payload['indexParams'])
|
|
# assert load success
|
|
assert rsp['data']['load'] == "LoadStateNotLoad"
|
|
|
|
@pytest.mark.parametrize("metric_type", ["L2", "IP", "COSINE"])
|
|
@pytest.mark.parametrize("dim", [128])
|
|
def test_create_collections_multi_float_vector_with_all_index(self, dim, metric_type):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"schema": {
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
|
|
{"fieldName": "book_describe", "dataType": "VarChar", "elementTypeParams": {"max_length": "256"}},
|
|
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}},
|
|
{"fieldName": "image_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}}
|
|
]
|
|
},
|
|
"indexParams": [
|
|
{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": f"{metric_type}"},
|
|
{"fieldName": "image_intro", "indexName": "image_intro_vector", "metricType": f"{metric_type}"}]
|
|
}
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
time.sleep(10)
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
# assert index created
|
|
indexes = rsp['data']['indexes']
|
|
assert len(indexes) == len(payload['indexParams'])
|
|
# assert load success
|
|
assert rsp['data']['load'] == "LoadStateLoaded"
|
|
|
|
@pytest.mark.parametrize("auto_id", [True])
|
|
@pytest.mark.parametrize("enable_dynamic_field", [True])
|
|
@pytest.mark.parametrize("enable_partition_key", [True])
|
|
@pytest.mark.parametrize("dim", [128])
|
|
@pytest.mark.parametrize("metric_type", ["L2", "IP", "COSINE"])
|
|
def test_create_collections_float16_vector_datatype(self, dim, auto_id, enable_dynamic_field, enable_partition_key,
|
|
metric_type):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"schema": {
|
|
"autoId": auto_id,
|
|
"enableDynamicField": enable_dynamic_field,
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "float16_vector", "dataType": "Float16Vector",
|
|
"elementTypeParams": {"dim": f"{dim}"}},
|
|
{"fieldName": "bfloat16_vector", "dataType": "BFloat16Vector",
|
|
"elementTypeParams": {"dim": f"{dim}"}},
|
|
]
|
|
},
|
|
"indexParams": [
|
|
{"fieldName": "float16_vector", "indexName": "float16_vector_index", "metricType": f"{metric_type}"},
|
|
{"fieldName": "bfloat16_vector", "indexName": "bfloat16_vector_index", "metricType": f"{metric_type}"}]
|
|
|
|
}
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
c = Collection(name)
|
|
logger.info(f"schema: {c.schema}")
|
|
# describe collection
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
assert len(rsp['data']['fields']) == len(c.schema.fields)
|
|
|
|
@pytest.mark.parametrize("auto_id", [True])
|
|
@pytest.mark.parametrize("enable_dynamic_field", [True])
|
|
@pytest.mark.parametrize("enable_partition_key", [True])
|
|
@pytest.mark.parametrize("dim", [128])
|
|
@pytest.mark.parametrize("metric_type", ["JACCARD", "HAMMING"])
|
|
@pytest.mark.skip(reason="https://github.com/milvus-io/milvus/issues/31494")
|
|
def test_create_collections_binary_vector_datatype(self, dim, auto_id, enable_dynamic_field, enable_partition_key,
|
|
metric_type):
|
|
"""
|
|
target: test create collection
|
|
method: create a collection with a simple schema
|
|
expected: create collection success
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"schema": {
|
|
"autoId": auto_id,
|
|
"enableDynamicField": enable_dynamic_field,
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "binary_vector", "dataType": "BinaryVector", "elementTypeParams": {"dim": f"{dim}"}},
|
|
]
|
|
},
|
|
"indexParams": [
|
|
{"fieldName": "binary_vector", "indexName": "binary_vector_index", "metricType": f"{metric_type}"}
|
|
]
|
|
|
|
}
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
c = Collection(name)
|
|
logger.info(f"schema: {c.schema}")
|
|
# describe collection
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
assert len(rsp['data']['fields']) == len(c.schema.fields)
|
|
|
|
def test_create_collections_concurrent_with_same_param(self):
|
|
"""
|
|
target: test create collection with same param
|
|
method: concurrent create collections with same param with multi thread
|
|
expected: create collections all success
|
|
"""
|
|
concurrent_rsp = []
|
|
|
|
def create_collection(c_name, vector_dim, c_metric_type):
|
|
collection_payload = {
|
|
"collectionName": c_name,
|
|
"dimension": vector_dim,
|
|
"metricType": c_metric_type,
|
|
}
|
|
rsp = client.collection_create(collection_payload)
|
|
concurrent_rsp.append(rsp)
|
|
logger.info(rsp)
|
|
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
metric_type = "L2"
|
|
client = self.collection_client
|
|
threads = []
|
|
for i in range(10):
|
|
t = threading.Thread(target=create_collection, args=(name, dim, metric_type,))
|
|
threads.append(t)
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
time.sleep(10)
|
|
success_cnt = 0
|
|
for rsp in concurrent_rsp:
|
|
if rsp['code'] == 0:
|
|
success_cnt += 1
|
|
logger.info(concurrent_rsp)
|
|
assert success_cnt == 10
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
|
|
def test_create_collections_concurrent_with_different_param(self):
|
|
"""
|
|
target: test create collection with different param
|
|
method: concurrent create collections with different param with multi thread
|
|
expected: only one collection can success
|
|
"""
|
|
concurrent_rsp = []
|
|
|
|
def create_collection(c_name, vector_dim, c_metric_type):
|
|
collection_payload = {
|
|
"collectionName": c_name,
|
|
"dimension": vector_dim,
|
|
"metricType": c_metric_type,
|
|
}
|
|
rsp = client.collection_create(collection_payload)
|
|
concurrent_rsp.append(rsp)
|
|
logger.info(rsp)
|
|
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
threads = []
|
|
for i in range(0, 5):
|
|
t = threading.Thread(target=create_collection, args=(name, dim + i, "L2",))
|
|
threads.append(t)
|
|
for i in range(5, 10):
|
|
t = threading.Thread(target=create_collection, args=(name, dim + i, "IP",))
|
|
threads.append(t)
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
time.sleep(10)
|
|
success_cnt = 0
|
|
for rsp in concurrent_rsp:
|
|
if rsp['code'] == 0:
|
|
success_cnt += 1
|
|
logger.info(concurrent_rsp)
|
|
assert success_cnt == 1
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
|
|
|
|
@pytest.mark.L1
|
|
class TestCreateCollectionNegative(TestBase):
|
|
|
|
def test_create_collections_custom_with_invalid_datatype(self):
|
|
"""
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"schema": {
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
|
|
{"fieldName": "book_describe", "dataType": "VARCHAR", "elementTypeParams": {"max_length": "256"}},
|
|
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}}
|
|
]
|
|
}
|
|
}
|
|
logging.info(f"create collection {name} with payload: {payload}")
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 1100
|
|
|
|
def test_create_collections_with_invalid_api_key(self):
|
|
"""
|
|
target: test create collection with invalid api key(wrong username and password)
|
|
method: create collections with invalid api key
|
|
expected: create collection failed
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
client.api_key = "illegal_api_key"
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 1800
|
|
|
|
@pytest.mark.parametrize("name",
|
|
[" ", "test_collection_" * 100, "test collection", "test/collection", "test\collection"])
|
|
def test_create_collections_with_invalid_collection_name(self, name):
|
|
"""
|
|
target: test create collection with invalid collection name
|
|
method: create collections with invalid collection name
|
|
expected: create collection failed with right error message
|
|
"""
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 1100
|
|
assert "Invalid collection name" in rsp['message'] or "invalid parameter" in rsp['message']
|
|
|
|
|
|
@pytest.mark.L0
|
|
class TestHasCollections(TestBase):
|
|
|
|
def test_has_collections_default(self):
|
|
"""
|
|
target: test list collection with a simple schema
|
|
method: create collections and list them
|
|
expected: created collections are in list
|
|
"""
|
|
client = self.collection_client
|
|
name_list = []
|
|
for i in range(2):
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
payload = {
|
|
"collectionName": name,
|
|
"metricType": "L2",
|
|
"dimension": dim,
|
|
}
|
|
time.sleep(1)
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
name_list.append(name)
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
for name in name_list:
|
|
assert name in all_collections
|
|
rsp = client.collection_has(collection_name=name)
|
|
assert rsp['data']['has'] is True
|
|
|
|
def test_has_collections_with_not_exist_name(self):
|
|
"""
|
|
target: test list collection with a simple schema
|
|
method: create collections and list them
|
|
expected: created collections are in list
|
|
"""
|
|
client = self.collection_client
|
|
name_list = []
|
|
for i in range(2):
|
|
name = gen_collection_name()
|
|
name_list.append(name)
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
for name in name_list:
|
|
assert name not in all_collections
|
|
rsp = client.collection_has(collection_name=name)
|
|
assert rsp['data']['has'] is False
|
|
|
|
|
|
@pytest.mark.L0
|
|
class TestGetCollectionStats(TestBase):
|
|
|
|
def test_get_collections_stats(self):
|
|
"""
|
|
target: test list collection with a simple schema
|
|
method: create collections and list them
|
|
expected: created collections are in list
|
|
"""
|
|
client = self.collection_client
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
payload = {
|
|
"collectionName": name,
|
|
"metricType": "L2",
|
|
"dimension": dim,
|
|
}
|
|
time.sleep(1)
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
# describe collection
|
|
client.collection_describe(collection_name=name)
|
|
rsp = client.collection_stats(collection_name=name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['rowCount'] == 0
|
|
# insert data
|
|
nb = 3000
|
|
data = get_data_by_payload(payload, nb)
|
|
payload = {
|
|
"collectionName": name,
|
|
"data": data
|
|
}
|
|
self.vector_client.vector_insert(payload=payload)
|
|
c = Collection(name)
|
|
count = c.query(expr="", output_fields=["count(*)"])
|
|
logger.info(f"count: {count}")
|
|
c.flush()
|
|
rsp = client.collection_stats(collection_name=name)
|
|
assert rsp['data']['rowCount'] == nb
|
|
|
|
|
|
class TestLoadReleaseCollection(TestBase):
|
|
|
|
def test_load_and_release_collection(self):
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"schema": {
|
|
"fields": [
|
|
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": True, "elementTypeParams": {}},
|
|
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
|
|
{"fieldName": "book_describe", "dataType": "VarChar", "elementTypeParams": {"max_length": "256"}},
|
|
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": f"{dim}"}}
|
|
]
|
|
}
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
# create index before load
|
|
index_params = [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]
|
|
payload = {
|
|
"collectionName": name,
|
|
"indexParams": index_params
|
|
}
|
|
rsp = self.index_client.index_create(payload)
|
|
|
|
# get load state before load
|
|
rsp = client.collection_load_state(collection_name=name)
|
|
assert rsp['data']['loadState'] == "LoadStateNotLoad"
|
|
|
|
# describe collection
|
|
client.collection_describe(collection_name=name)
|
|
rsp = client.collection_load(collection_name=name)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_load_state(collection_name=name)
|
|
assert rsp['data']['loadState'] in ["LoadStateLoaded", "LoadStateLoading"]
|
|
time.sleep(5)
|
|
rsp = client.collection_load_state(collection_name=name)
|
|
assert rsp['data']['loadState'] == "LoadStateLoaded"
|
|
|
|
# release collection
|
|
rsp = client.collection_release(collection_name=name)
|
|
time.sleep(5)
|
|
rsp = client.collection_load_state(collection_name=name)
|
|
assert rsp['data']['loadState'] == "LoadStateNotLoad"
|
|
|
|
@pytest.mark.L0
|
|
class TestGetCollectionLoadState(TestBase):
|
|
|
|
def test_get_collection_load_state(self):
|
|
"""
|
|
target: test list collection with a simple schema
|
|
method: create collections and list them
|
|
expected: created collections are in list
|
|
"""
|
|
client = self.collection_client
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
payload = {
|
|
"collectionName": name,
|
|
"metricType": "L2",
|
|
"dimension": dim,
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
# describe collection
|
|
client.collection_describe(collection_name=name)
|
|
rsp = client.collection_load_state(collection_name=name)
|
|
assert rsp['code'] == 0
|
|
t0 = time.time()
|
|
while time.time() - t0 < 10:
|
|
rsp = client.collection_load_state(collection_name=name)
|
|
if rsp['data']['loadState'] != "LoadStateNotLoad":
|
|
break
|
|
time.sleep(1)
|
|
assert rsp['data']['loadState'] in ["LoadStateLoading", "LoadStateLoaded"]
|
|
# insert data
|
|
nb = 3000
|
|
data = get_data_by_payload(payload, nb)
|
|
payload = {
|
|
"collectionName": name,
|
|
"data": data
|
|
}
|
|
self.vector_client.vector_insert(payload=payload)
|
|
rsp = client.collection_load_state(collection_name=name)
|
|
assert rsp['data']['loadState'] in ["LoadStateLoading", "LoadStateLoaded"]
|
|
time.sleep(10)
|
|
rsp = client.collection_load_state(collection_name=name)
|
|
assert rsp['data']['loadState'] == "LoadStateLoaded"
|
|
|
|
|
|
@pytest.mark.L0
|
|
class TestListCollections(TestBase):
|
|
|
|
def test_list_collections_default(self):
|
|
"""
|
|
target: test list collection with a simple schema
|
|
method: create collections and list them
|
|
expected: created collections are in list
|
|
"""
|
|
client = self.collection_client
|
|
name_list = []
|
|
for i in range(2):
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
payload = {
|
|
"collectionName": name,
|
|
"metricType": "L2",
|
|
"dimension": dim,
|
|
}
|
|
time.sleep(1)
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
name_list.append(name)
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
for name in name_list:
|
|
assert name in all_collections
|
|
|
|
|
|
@pytest.mark.L1
|
|
class TestListCollectionsNegative(TestBase):
|
|
def test_list_collections_with_invalid_api_key(self):
|
|
"""
|
|
target: test list collection with an invalid api key
|
|
method: list collection with invalid api key
|
|
expected: raise error with right error code and message
|
|
"""
|
|
client = self.collection_client
|
|
name_list = []
|
|
for i in range(2):
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
payload = {
|
|
"collectionName": name,
|
|
"metricType": "L2",
|
|
"dimension": dim,
|
|
}
|
|
time.sleep(1)
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
name_list.append(name)
|
|
client = self.collection_client
|
|
client.api_key = "illegal_api_key"
|
|
rsp = client.collection_list()
|
|
assert rsp['code'] == 1800
|
|
|
|
|
|
@pytest.mark.L0
|
|
class TestDescribeCollection(TestBase):
|
|
|
|
def test_describe_collections_default(self):
|
|
"""
|
|
target: test describe collection with a simple schema
|
|
method: describe collection
|
|
expected: info of description is same with param passed to create collection
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
"metricType": "L2"
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
assert rsp['data']['autoId'] is False
|
|
assert rsp['data']['enableDynamicField'] is True
|
|
assert len(rsp['data']['indexes']) == 1
|
|
|
|
def test_describe_collections_custom(self):
|
|
"""
|
|
target: test describe collection with a simple schema
|
|
method: describe collection
|
|
expected: info of description is same with param passed to create collection
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
fields = [
|
|
FieldSchema(name='reviewer_id', dtype=DataType.INT64, description="", is_primary=True),
|
|
FieldSchema(name='store_address', dtype=DataType.VARCHAR, description="", max_length=512,
|
|
is_partition_key=True),
|
|
FieldSchema(name='review', dtype=DataType.VARCHAR, description="", max_length=16384),
|
|
FieldSchema(name='vector', dtype=DataType.FLOAT_VECTOR, description="", dim=384, is_index=True),
|
|
]
|
|
|
|
schema = CollectionSchema(
|
|
fields=fields,
|
|
description="",
|
|
enable_dynamic_field=True,
|
|
# The following is an alternative to setting `is_partition_key` in a field schema.
|
|
partition_key_field="store_address"
|
|
)
|
|
|
|
collection = Collection(
|
|
name=name,
|
|
schema=schema,
|
|
)
|
|
logger.info(f"schema: {schema}")
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
rsp = client.collection_describe(name)
|
|
assert rsp['code'] == 0
|
|
assert rsp['data']['collectionName'] == name
|
|
|
|
for field in rsp['data']['fields']:
|
|
if field['name'] == "store_address":
|
|
assert field['partitionKey'] is True
|
|
if field['name'] == "reviewer_id":
|
|
assert field['primaryKey'] is True
|
|
assert rsp['data']['autoId'] is False
|
|
assert rsp['data']['enableDynamicField'] is True
|
|
|
|
|
|
@pytest.mark.L1
|
|
class TestDescribeCollectionNegative(TestBase):
|
|
def test_describe_collections_with_invalid_api_key(self):
|
|
"""
|
|
target: test describe collection with invalid api key
|
|
method: describe collection with invalid api key
|
|
expected: raise error with right error code and message
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
illegal_client = CollectionClient(self.endpoint, "illegal_api_key")
|
|
rsp = illegal_client.collection_describe(name)
|
|
assert rsp['code'] == 1800
|
|
|
|
def test_describe_collections_with_invalid_collection_name(self):
|
|
"""
|
|
target: test describe collection with invalid collection name
|
|
method: describe collection with invalid collection name
|
|
expected: raise error with right error code and message
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# describe collection
|
|
invalid_name = "invalid_name"
|
|
rsp = client.collection_describe(invalid_name)
|
|
assert rsp['code'] == 100
|
|
assert "can't find collection" in rsp['message']
|
|
|
|
|
|
@pytest.mark.L0
|
|
class TestDropCollection(TestBase):
|
|
def test_drop_collections_default(self):
|
|
"""
|
|
Drop a collection with a simple schema
|
|
target: test drop collection with a simple schema
|
|
method: drop collection
|
|
expected: dropped collection was not in collection list
|
|
"""
|
|
clo_list = []
|
|
for i in range(5):
|
|
time.sleep(1)
|
|
name = 'test_collection_' + datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S_%f_%f")
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": 128,
|
|
"metricType": "L2"
|
|
}
|
|
rsp = self.collection_client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
clo_list.append(name)
|
|
rsp = self.collection_client.collection_list()
|
|
all_collections = rsp['data']
|
|
for name in clo_list:
|
|
assert name in all_collections
|
|
for name in clo_list:
|
|
time.sleep(0.2)
|
|
payload = {
|
|
"collectionName": name,
|
|
}
|
|
rsp = self.collection_client.collection_drop(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = self.collection_client.collection_list()
|
|
all_collections = rsp['data']
|
|
for name in clo_list:
|
|
assert name not in all_collections
|
|
|
|
|
|
@pytest.mark.L1
|
|
class TestDropCollectionNegative(TestBase):
|
|
def test_drop_collections_with_invalid_api_key(self):
|
|
"""
|
|
target: test drop collection with invalid api key
|
|
method: drop collection with invalid api key
|
|
expected: raise error with right error code and message; collection still in collection list
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# drop collection
|
|
payload = {
|
|
"collectionName": name,
|
|
}
|
|
illegal_client = CollectionClient(self.endpoint, "invalid_api_key")
|
|
rsp = illegal_client.collection_drop(payload)
|
|
assert rsp['code'] == 1800
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
|
|
def test_drop_collections_with_invalid_collection_name(self):
|
|
"""
|
|
target: test drop collection with invalid collection name
|
|
method: drop collection with invalid collection name
|
|
expected: raise error with right error code and message
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"dimension": dim,
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
# drop collection
|
|
invalid_name = "invalid_name"
|
|
payload = {
|
|
"collectionName": invalid_name,
|
|
}
|
|
rsp = client.collection_drop(payload)
|
|
assert rsp['code'] == 0
|
|
|
|
|
|
@pytest.mark.L0
|
|
class TestRenameCollection(TestBase):
|
|
|
|
def test_rename_collection(self):
|
|
"""
|
|
target: test rename collection
|
|
method: rename collection
|
|
expected: renamed collection is in collection list
|
|
"""
|
|
name = gen_collection_name()
|
|
dim = 128
|
|
client = self.collection_client
|
|
payload = {
|
|
"collectionName": name,
|
|
"metricType": "L2",
|
|
"dimension": dim,
|
|
}
|
|
rsp = client.collection_create(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert name in all_collections
|
|
new_name = gen_collection_name()
|
|
payload = {
|
|
"collectionName": name,
|
|
"newCollectionName": new_name,
|
|
}
|
|
rsp = client.collection_rename(payload)
|
|
assert rsp['code'] == 0
|
|
rsp = client.collection_list()
|
|
all_collections = rsp['data']
|
|
assert new_name in all_collections
|
|
assert name not in all_collections
|