fix upsert/load cases

Signed-off-by: ThreadDao <[email protected]>
milvus-io · Jun 28, 2024 · a7dc18b · a7dc18b
1 parent b62fdb2
commit a7dc18b
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 14 deletions.
diff --git a/test/testcases/load_release_test.go b/test/testcases/load_release_test.go
@@ -1,4 +1,4 @@
-//go:build L0
+///go:build L0
 
 package testcases
 
@@ -247,16 +247,24 @@ func TestLoadMultiPartitions(t *testing.T) {
 	idx, _ := entity.NewIndexHNSW(entity.L2, 8, 96)
 	mc.CreateIndex(ctx, collName, common.DefaultFloatVecFieldName, idx, false)
 
-	// load partition
-	errLoad := mc.LoadPartitions(ctx, collName, []string{partitionName, common.DefaultPartition}, false)
+	// load default partition
+	errLoad := mc.LoadPartitions(ctx, collName, []string{common.DefaultPartition}, false)
 	common.CheckErr(t, errLoad, true)
 
-	//query nb from partition
-	queryIds := entity.NewColumnInt64(common.DefaultIntFieldName, []int64{0, common.DefaultNb})
-	queryResultPartition, _ := mc.QueryByPks(ctx, collName, []string{}, queryIds, []string{common.DefaultIntFieldName})
-	common.CheckQueryResult(t, queryResultPartition, []entity.Column{
-		entity.NewColumnInt64(common.DefaultIntFieldName, []int64{0, common.DefaultNb}),
-	})
+	//query nb from default partition
+	resDef, _ := mc.Query(ctx, collName, []string{common.DefaultPartition}, "", []string{common.QueryCountFieldName})
+	require.EqualValues(t, common.DefaultNb, resDef.GetColumn(common.QueryCountFieldName).(*entity.ColumnInt64).Data()[0])
+
+	// load partition and query -> actually not loaded
+	errLoad = mc.LoadPartitions(ctx, collName, []string{partitionName}, false)
+	common.CheckErr(t, errLoad, true)
+	resPar, _ := mc.Query(ctx, collName, []string{partitionName}, "", []string{common.QueryCountFieldName})
+	require.EqualValues(t, 0, resPar.GetColumn(common.QueryCountFieldName).(*entity.ColumnInt64).Data()[0])
+
+	mc.ReleaseCollection(ctx, collName)
+	mc.LoadCollection(ctx, collName, false)
+	res, _ := mc.Query(ctx, collName, []string{}, "", []string{common.QueryCountFieldName})
+	require.EqualValues(t, common.DefaultNb*2, res.GetColumn(common.QueryCountFieldName).(*entity.ColumnInt64).Data()[0])
 }
 
 // test load partitions repeatedly

diff --git a/test/testcases/upsert_test.go b/test/testcases/upsert_test.go
@@ -1,4 +1,4 @@
-//go:build L0
+///go:build L0
 
 package testcases
 
@@ -219,12 +219,15 @@ func TestUpsertNotExistCollectionPartition(t *testing.T) {
 	mc := createMilvusClient(ctx, t)
 
 	// create default collection with autoID true
-	collName := createDefaultCollection(ctx, t, mc, true, common.DefaultShards)
+	collName := createDefaultCollection(ctx, t, mc, false, common.DefaultShards)
 
 	// upsert not exist partition
-	_, floatColumn, vecColumn := common.GenDefaultColumnData(0, common.DefaultNb, common.DefaultDim)
+	pkColumn, floatColumn, vecColumn := common.GenDefaultColumnData(0, common.DefaultNb, common.DefaultDim)
 	_, errUpsert := mc.Upsert(ctx, collName, "aaa", floatColumn, vecColumn)
-	common.CheckErr(t, errUpsert, false, "can not assign primary field data when auto id enabled")
+	common.CheckErr(t, errUpsert, false, "field int64 not passed")
+
+	_, errUpsert = mc.Upsert(ctx, collName, "aaa", pkColumn, floatColumn, vecColumn)
+	common.CheckErr(t, errUpsert, false, "partition not found[partition=aaa]")
 
 	// upsert not exist collection
 	_, errUpsert = mc.Upsert(ctx, "aaa", "", floatColumn, vecColumn)
@@ -382,7 +385,9 @@ func TestUpsertPartitionKeyCollection(t *testing.T) {
 
 	// upsert data partition key field [nb, nb*2)
 	partitionKeyColumn2 := common.GenColumnData(common.DefaultNb, common.DefaultNb, entity.FieldTypeInt64, partitionKeyFieldName)
-	mc.Upsert(ctx, schema.CollectionName, "", pkColumn, floatColumn, vecColumn, partitionKeyColumn2)
+	_, _, vecColumn2 := common.GenDefaultColumnData(0, common.DefaultNb, common.DefaultDim)
+	_, err = mc.Upsert(ctx, schema.CollectionName, "", pkColumn, floatColumn, vecColumn2, partitionKeyColumn2)
+	common.CheckErr(t, err, true)
 
 	// verify upsert
 	resSet, err := mc.Query(ctx, schema.CollectionName, []string{}, fmt.Sprintf("%s >= %d", partitionKeyFieldName, common.DefaultNb),
@@ -392,6 +397,10 @@ func TestUpsertPartitionKeyCollection(t *testing.T) {
 	resSet, err = mc.Query(ctx, schema.CollectionName, []string{}, fmt.Sprintf("%s < %d", partitionKeyFieldName, common.DefaultNb),
 		[]string{common.QueryCountFieldName})
 	require.Equal(t, int64(0), resSet.GetColumn(common.QueryCountFieldName).(*entity.ColumnInt64).Data()[0])
+
+	resSet, err = mc.Query(ctx, schema.CollectionName, []string{}, fmt.Sprintf("%s < %d", common.DefaultIntFieldName, 10),
+		[]string{"*"})
+	common.CheckQueryResult(t, resSet, []entity.Column{pkColumn.Slice(0, 10), floatColumn.Slice(0, 10), partitionKeyColumn2.Slice(0, 10), vecColumn2.Slice(0, 10)})
 }
 
 func TestUpsertWithoutLoading(t *testing.T) {