Skip to content

Commit

Permalink
test: add index check after restore
Browse files Browse the repository at this point in the history 
Signed-off-by: zhuwenxing <[email protected]>
  • Loading branch information
@zhuwenxing
zhuwenxing committed Jan 16, 2025
1 parent 312e7e9 commit cee1562
Show file tree
Hide file tree
Showing 2 changed files with 129 additions and 17 deletions.
42 changes: 25 additions & 17 deletions tests/base/client_base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -513,8 +513,15 @@ def is_binary_by_schema(self, schema):
if field.dtype == DataType.FLOAT_VECTOR:
return False

def compare_indexes(self, src_name, dist_name):
collection_src, _ = self.collection_wrap.init_collection(name=src_name)
collection_dist, _ = self.collection_wrap.init_collection(name=dist_name)
log.info(f"collection_src schema: {collection_src.schema}")
log.info(f"collection_dist schema: {collection_dist.schema}")
assert collection_src.indexes == collection_dist.indexes

def compare_collections(
self, src_name, dist_name, output_fields=None, verify_by_query=False
self, src_name, dist_name, output_fields=None, verify_by_query=False, skip_index=False
):
if output_fields is None:
output_fields = ["*"]
Expand All @@ -538,22 +545,23 @@ def compare_collections(
assert src_num == dist_num, f"srs_num: {src_num}, dist_num: {dist_num}"
return
for coll in [collection_src, collection_dist]:
is_binary = self.is_binary_by_schema(coll.schema)
try:
if is_binary:
coll.create_index(
ct.default_binary_vec_field_name,
ct.default_bin_flat_index,
index_name=cf.gen_unique_str(),
)
else:
coll.create_index(
ct.default_float_vec_field_name,
ct.default_index,
index_name=cf.gen_unique_str(),
)
except Exception as e:
log.error(f"collection {coll.name} create index failed with error: {e}")
if not skip_index:
is_binary = self.is_binary_by_schema(coll.schema)
try:
if is_binary:
coll.create_index(
ct.default_binary_vec_field_name,
ct.default_bin_flat_index,
index_name=cf.gen_unique_str(),
)
else:
coll.create_index(
ct.default_float_vec_field_name,
ct.default_index,
index_name=cf.gen_unique_str(),
)
except Exception as e:
log.error(f"collection {coll.name} create index failed with error: {e}")
coll.load()
time.sleep(5)
# get entities by count
Expand Down
104 changes: 104 additions & 0 deletions tests/testcases/test_restore_backup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,110 @@ def test_milvus_restore_back(
for name in restore_collections:
self.compare_collections(name, name + suffix, verify_by_query=True)

@pytest.mark.parametrize("nb", [3000])
@pytest.mark.parametrize("is_auto_id", [True])
@pytest.mark.parametrize("enable_partition", [False])
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("collection_need_to_restore", [3])
@pytest.mark.parametrize("collection_type", ["all"])
@pytest.mark.tags(CaseLabel.L0)
def test_milvus_restore_back_with_index(
self,
collection_type,
collection_need_to_restore,
is_async,
is_auto_id,
enable_partition,
nb,
):
# prepare data
names_origin = []
back_up_name = cf.gen_unique_str(backup_prefix)
if collection_type == "all":
for is_binary in [True, False, False]:
names_origin.append(cf.gen_unique_str(prefix))
self.prepare_data(
names_origin[-1],
nb=nb,
is_binary=is_binary,
auto_id=is_auto_id,
check_function=False,
enable_partition=enable_partition,
)
if collection_type == "float":
for is_binary in [False, False, False]:
names_origin.append(cf.gen_unique_str(prefix))
self.prepare_data(
names_origin[-1],
nb=nb,
is_binary=is_binary,
auto_id=is_auto_id,
check_function=False,
enable_partition=enable_partition,
)
if collection_type == "binary":
for is_binary in [True, True, True]:
names_origin.append(cf.gen_unique_str(prefix))
self.prepare_data(
names_origin[-1],
nb=nb,
is_binary=is_binary,
auto_id=is_auto_id,
check_function=False,
enable_partition=enable_partition,
)
log.info(f"name_origin:{names_origin}, back_up_name: {back_up_name}")
for name in names_origin:
res, _ = self.utility_wrap.has_collection(name)
assert res is True
# create backup
names_need_backup = names_origin
payload = {
"async": False,
"backup_name": back_up_name,
"collection_names": names_need_backup,
}
res = self.client.create_backup(payload)
log.info(f"create backup response: {res}")
backup = self.client.get_backup(back_up_name)
assert backup["data"]["name"] == back_up_name
backup_collections = [
backup["collection_name"] for backup in backup["data"]["collection_backups"]
]
restore_collections = backup_collections
if collection_need_to_restore == "all":
payload = {
"async": False,
"backup_name": back_up_name,
"collection_suffix": suffix,
}
else:
restore_collections = names_need_backup[:collection_need_to_restore]
payload = {
"async": False,
"backup_name": back_up_name,
"collection_suffix": suffix,
"collection_names": restore_collections,
}
payload["restoreIndex"] = True
t0 = time.time()
res = self.client.restore_backup(payload)
restore_id = res["data"]["id"]
log.info(f"restore_backup: {res}")
if is_async:
res = self.client.wait_restore_complete(restore_id)
assert res is True
t1 = time.time()
log.info(f"restore {restore_collections} cost time: {t1 - t0}")
res, _ = self.utility_wrap.list_collections()
for name in restore_collections:
assert name + suffix in res
for name in restore_collections:
self.compare_index(name, name + suffix, verify_by_query=False)
for name in restore_collections:
self.compare_collections(name, name + suffix, verify_by_query=True, skip_index=True)


@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize("nb", [3000])
@pytest.mark.parametrize("is_auto_id", [True])
Expand Down

0 comments on commit cee1562

Please sign in to comment.