def test_false_vector(self, gcon):
param = {
'table_name': fake.table_name(),
'query_records': records_factory(dim, nq),
'top_k': 'string',
'nprobe': 10
}
with pytest.raises(ParamError):
gcon.search(**param)
param = {
'table_name': fake.table_name(),
'query_records': records_factory(dim, nq),
'top_k': 'string',
'nprobe': 10
}
with pytest.raises(ParamError):
gcon.search(**param)
param = {
'table_name': fake.table_name(),
'query_records': records_factory(dim, nq),
'top_k': random.randint(1, 10),
'nprobe': 10,
'query_ranges': ['false_date_format']
}
with pytest.raises(ParamError):
gcon.search(**param)
def test_insert_with_mismatch_field(self, connect, hvcollection, dim):
entities = [{"Vec": records_factory(dim, 1)[0], "Int": 1, "Float": 0.2}]
with pytest.raises(ParamError):
connect.insert(hvcollection, entities)
entities = [{"Vec": records_factory(dim, 1)[0]}]
with pytest.raises(ParamError):
connect.insert(hvcollection, entities)
def test_false_insert(self, gcon):
param = {
'table_name': fake.table_name(),
'records': records_factory(dim, nq)
}
res, ids = gcon.insert(**param)
assert not res.OK()
def test_insert_with_wrong_type_data(self, connect, hvcollection, dim):
vectors = records_factory(dim, 10000)
integers = integer_factory(10000)
entities = [{"Vec": v, "Int": i} for v, i in zip(vectors, integers)]
connect.insert(hvcollection, entities)
def test_insert_wrong_table_name(self, gcon):
table_name = "&*^%&&dvfdgd(()"
vectors = records_factory(dim, nq)
status, _ = gcon.insert(table_name, vectors)
assert not status.OK()
def test_search_multi_condition_in_single_must(self, connect, ivrecords, dim):
topk = 10
xq = records_factory(dim, 10)
xi = integer_factory(2)
metric = "L2"
dsl = {
"bool": {
"must": [
{
"term": {"Int": xi}
},
{
"vector": {
"Vec": {"topk": topk, "query": xq, "metric_type": metric}
}
}
]
}
}
try:
connect.search(ivrecords, dsl)
except Exception as e:
pytest.fail(f"{e}")
def test_insert_scalar_with_id(self, scalar, vec, connect, dim):
collection_name = fake.collection_name()
collection_param = {
"fields": [
{"name": "scalar", "type": scalar},
{"name": "Vec", "type": vec, "params": {"dim": 128}}
],
"auto_id": False
}
try:
connect.create_collection(collection_name, collection_param)
except Exception as e:
pytest.fail(f"Create collection {collection_name} fail: {e}")
if scalar in (DataType.INT32, DataType.INT64):
scalars = [random.randint(0, 10000) for _ in range(10000)]
else:
scalars = [random.random() for _ in range(10000)]
if vec in (DataType.FLOAT_VECTOR,):
vectors = records_factory(dim, 10000)
else:
vectors = binary_records_factory(dim, 10000)
entities = [{"Vec": vector, "scalar": s} for vector, s in zip(vectors, scalars)]
for i, e in enumerate(entities):
e["_id"] = i
try:
connect.insert(collection_name, entities)
except Exception as e:
pytest.fail(f"Unexpected Insert Error: {e}")
finally:
connect.drop_collection(collection_name)
def test_bulk_insert_async_with_collection_non_exist(
self, connect, hvcollection, dim):
collection = hvcollection + "_c1"
vectors = records_factory(dim, 1)
integers = integer_factory(1)
entities = [{
"name": "Vec",
"values": vectors
}, {
"name": "Int",
"values": integers
}]
with pytest.raises(BaseError):
future = connect.bulk_insert(collection, entities, _async=True)
future.result()
entities = [{
"name": "Vec",
"values": vectors,
"type": DataType.FLOAT_VECTOR
}, {
"name": "Int",
"values": integers,
"type": DataType.INT64
}]
future = connect.bulk_insert(collection, entities, _async=True)
with pytest.raises(BaseError):
future.result()
def test_create_collection_exception(self, connect, hvcollection, dim):
vectors = records_factory(dim, 1)
integers = integer_factory(1)
entities = [{
"name": "Vec",
"values": vectors
}, {
"name": "Int",
"values": integers
}]
mock_grpc_timeout = mock.MagicMock(
side_effect=grpc.FutureTimeoutError())
with mock.patch.object(Uum, 'future', mock_grpc_timeout):
with pytest.raises(grpc.FutureTimeoutError):
connect.bulk_insert(hvcollection, entities)
mock_grpc_error = mock.MagicMock(side_effect=MockGrpcError())
with mock.patch.object(Uum, 'future', mock_grpc_error):
with pytest.raises(grpc.RpcError):
connect.bulk_insert(hvcollection, entities)
mock_exception = mock.MagicMock(side_effect=Exception("error"))
with mock.patch.object(Uum, 'future', mock_exception):
with pytest.raises(Exception):
connect.bulk_insert(hvcollection, entities)
def test_bulk_insert_with_wrong_scalar_type_then_correct_type(
self, connect, hvcollection, dim):
vectors = records_factory(dim, 10000)
integers = integer_factory(10000)
entities = [{
"name": "Vec",
"values": vectors,
"type": DataType.FLOAT_VECTOR
}, {
"name": "Int",
"values": integers,
"type": DataType.FLOAT
}]
with pytest.raises(BaseError):
connect.bulk_insert(hvcollection, entities)
entities = [{
"name": "Vec",
"values": vectors,
"type": DataType.FLOAT_VECTOR
}, {
"name": "Int",
"values": integers,
"type": DataType.INT64
}]
try:
connect.bulk_insert(hvcollection, entities)
except Exception as e:
pytest.fail("Unexpected MyError: {}".format(str(e)))
def test_insert(self, gcon, gtable):
param = {'table_name': gtable, 'records': records_factory(dim, nq)}
res, ids = gcon.insert(**param)
assert res.OK()
assert isinstance(ids, list)
assert len(ids) == nq
def test_search_scalar_flat(self, connect, ivrecords, dim):
query_vectors = records_factory(dim, 10)
dsl = vector_dsl(10, query_vectors)
results = connect.search(ivrecords, dsl)
assert len(results) == 10
assert len(results[0]) == 10
def test_search_flat(self, metric, connect, vrecords, dim):
query_vectors = records_factory(dim, 10)
dsl = vector_dsl(10, query_vectors, metric)
results = connect.search(vrecords, dsl)
assert len(results) == 10
assert len(results[0]) == 10
def test_flush_async_callback(self, gcon, gcollection):
def cb(status):
assert status.OK()
records = records_factory(dim, nq)
gcon.insert(gcollection, records)
future = gcon.flush([gcollection], _async=True, _callback=cb)
future.done()
def test_insert_with_no_right_dimension(self, gcon, gcollection):
param = {
'collection_name': gcollection,
'records': records_factory(dim + 1, nq)
}
res, ids = gcon.insert(**param)
assert not res.OK()
def test_insert_normal(self, connect, vcollection, dim):
vectors = records_factory(dim, 10000)
entities = [{"Vec": vector} for vector in vectors]
try:
connect.insert(vcollection, entities)
except Exception as e:
pytest.fail(f"Unexpected MyError: {e}")
def test_delete_entity_by_id_normal(self, gcon, gcollection):
vectors = records_factory(dim, nq)
status, ids = gcon.insert(gcollection, vectors)
gcon.flush([gcollection])
assert status.OK()
status = gcon.delete_entity_by_id(gcollection, ids[0:10])
assert status.OK()
def test_bulk_insert_normal(self, connect, vcollection, dim):
vectors = records_factory(dim, 10000)
entities = [{"name": "Vec", "values": vectors}]
try:
connect.bulk_insert(vcollection, entities)
except Exception as e:
pytest.fail("Unexpected MyError: {}".format(str(e)))
def test_add_vector_with_no_right_dimension(self, gcon, gtable):
param = {
'table_name': gtable,
'records': records_factory(dim + 1, nq)
}
res, ids = gcon.add_vectors(**param)
assert not res.OK()
def test_delete_by_id_normal(self, gcon, gcollection):
vectors = records_factory(dim, nq)
status, ids = gcon.insert(gcollection, vectors)
assert status.OK()
time.sleep(2)
status = gcon.delete_by_id(gcollection, ids[0:10])
assert status.OK()
def test_bulk_insert_with_id(self, connect, vicollection, dim):
vectors = records_factory(dim, 10000)
entities = [{"name": "Vec", "values": vectors}]
ids = [i for i in range(10000)]
try:
connect.bulk_insert(vicollection, entities, ids)
except Exception as e:
pytest.fail("Unexpected MyError: {}".format(str(e)))
def test_search_by_id_with_exceed_id(self, gcon, gcollection):
vectors = records_factory(dim, nq)
status, ids = gcon.insert(gcollection, vectors)
assert status.OK()
status, result = gcon.search_by_id(gcollection, 2, 10, ids[0] + 100)
assert status.OK()
print(result)
assert 0 == len(result)
def test_insert_with_wrong_nonaligned_data(self, connect, hvcollection, dim):
vectors = records_factory(dim, 10000)
integers = integer_factory(9998)
entities = [{"Vec": v, "Int": i} for v, i in zip(vectors[:9998], integers[:9998])]
entities.append({"Vec": vectors[9998]})
with pytest.raises(BaseError):
connect.insert(hvcollection, entities)
def test_search_in_files_wrong_file_ids(self, gcon, gvector):
param = {
'table_name': gvector,
'query_records': records_factory(dim, nq),
'file_ids': ['3388833'],
'top_k': random.randint(1, 10)
}
sta, results = gcon.search_in_files(**param)
assert not sta.OK()
def gvector(request, gcon, gcollection):
dim = getattr(request.module, 'dim', 128)
records = records_factory(dim, 10000)
gcon.insert(gcollection, records)
gcon.flush([gcollection])
return gcollection
def test_search_with_empty_collection(self, connect, vcollection, dim):
query_vectors = records_factory(dim, 10)
dsl = vector_dsl(10, query_vectors)
results = connect.search(vcollection, dsl)
assert len(results) == 10
for result in results:
assert len(result) == 0
def test_search_with_fields(self, connect, ivrecords, dim):
xq = records_factory(dim, 10)
dsl = vector_dsl(10, xq)
results = connect.search(ivrecords, dsl, fields=["Int", "Vec"])
assert len(results) == 10
assert len(results[0]) == 10
assert isinstance(results[0][0].entity.get("Int"), int)
assert len(results[0][0].entity.get("Vec")) == dim
def gvector(request, gcon, gtable):
dim = getattr(request.module, 'dim', 128)
records = records_factory(dim, 10000)
gcon.insert(gtable, records)
time.sleep(3)
return gtable
def test_insert_with_wrong_ids(self, gcon, gtable):
param = {
'table_name': gtable,
'records': records_factory(dim, nq),
'ids': [i + 1 for i in range(nq - 3)]
}
with pytest.raises(ParamError):
gcon.insert(**param)
def test_flush_async_normal(self, gcon, gcollection, ghandler):
if ghandler == "HTTP":
pytest.skip("HTTP handler not support async")
records = records_factory(dim, nq)
gcon.insert(gcollection, records)
future = gcon.flush([gcollection], _async=True)
status = future.result()
assert status.OK()