enhance: unify data generation test APIs (#32955)

Issue: #22837 Signed-off-by: Cai Yudong <[email protected]>
milvus-io · May 14, 2024 · 4fc7915 · 4fc7915
1 parent 96489b8
commit 4fc7915
Show file tree

Hide file tree

Showing 18 changed files with 1,092 additions and 1,464 deletions.
diff --git a/internal/proxy/mock_test.go b/internal/proxy/mock_test.go
diff --git a/internal/proxy/msg_pack_test.go b/internal/proxy/msg_pack_test.go
@@ -32,11 +32,12 @@ import (
 	"github.com/milvus-io/milvus/pkg/mq/msgstream"
 	"github.com/milvus-io/milvus/pkg/util/funcutil"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/testutils"
 )
 
 func TestRepackInsertData(t *testing.T) {
 	nb := 10
-	hash := generateHashKeys(nb)
+	hash := testutils.GenerateHashKeys(nb)
 	prefix := "TestRepackInsertData"
 	dbName := ""
 	collectionName := prefix + funcutil.GenRandomStr()
@@ -143,7 +144,7 @@ func TestRepackInsertData(t *testing.T) {
 
 func TestRepackInsertDataWithPartitionKey(t *testing.T) {
 	nb := 10
-	hash := generateHashKeys(nb)
+	hash := testutils.GenerateHashKeys(nb)
 	prefix := "TestRepackInsertData"
 	collectionName := prefix + funcutil.GenRandomStr()
 

diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go
@@ -73,6 +73,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/metric"
 	"github.com/milvus-io/milvus/pkg/util/metricsinfo"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/testutils"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
@@ -599,7 +600,7 @@ func TestProxy(t *testing.T) {
 	constructCollectionInsertRequest := func() *milvuspb.InsertRequest {
 		fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
 		bVecColumn := newBinaryVectorFieldData(binaryVecField, rowNum, dim)
-		hashKeys := generateHashKeys(rowNum)
+		hashKeys := testutils.GenerateHashKeys(rowNum)
 		return &milvuspb.InsertRequest{
 			Base:           nil,
 			DbName:         dbName,
@@ -614,7 +615,7 @@ func TestProxy(t *testing.T) {
 	constructPartitionInsertRequest := func() *milvuspb.InsertRequest {
 		fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
 		bVecColumn := newBinaryVectorFieldData(binaryVecField, rowNum, dim)
-		hashKeys := generateHashKeys(rowNum)
+		hashKeys := testutils.GenerateHashKeys(rowNum)
 		return &milvuspb.InsertRequest{
 			Base:           nil,
 			DbName:         dbName,
@@ -629,7 +630,7 @@ func TestProxy(t *testing.T) {
 	constructCollectionUpsertRequest := func() *milvuspb.UpsertRequest {
 		fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
 		bVecColumn := newBinaryVectorFieldData(binaryVecField, rowNum, dim)
-		hashKeys := generateHashKeys(rowNum)
+		hashKeys := testutils.GenerateHashKeys(rowNum)
 		return &milvuspb.UpsertRequest{
 			Base:           nil,
 			DbName:         dbName,
@@ -1811,7 +1812,7 @@ func TestProxy(t *testing.T) {
 					Dim: int64(dim),
 					Data: &schemapb.VectorField_FloatVector{
 						FloatVector: &schemapb.FloatArray{
-							Data: generateFloatVectors(nq, dim),
+							Data: testutils.GenerateFloatVectors(nq, dim),
 						},
 					},
 				},
@@ -1824,7 +1825,7 @@ func TestProxy(t *testing.T) {
 					Dim: int64(dim),
 					Data: &schemapb.VectorField_FloatVector{
 						FloatVector: &schemapb.FloatArray{
-							Data: generateFloatVectors(nq, dim),
+							Data: testutils.GenerateFloatVectors(nq, dim),
 						},
 					},
 				},
@@ -3723,7 +3724,7 @@ func TestProxy(t *testing.T) {
 		pkFieldData := newScalarFieldData(schema.Fields[0], int64Field, rowNum)
 		fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
 		bVecColumn := newBinaryVectorFieldData(binaryVecField, rowNum, dim)
-		hashKeys := generateHashKeys(rowNum)
+		hashKeys := testutils.GenerateHashKeys(rowNum)
 		return &milvuspb.UpsertRequest{
 			Base:           nil,
 			DbName:         dbName,
@@ -3739,7 +3740,7 @@ func TestProxy(t *testing.T) {
 		pkFieldData := newScalarFieldData(schema.Fields[0], int64Field, rowNum)
 		fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
 		bVecColumn := newBinaryVectorFieldData(binaryVecField, rowNum, dim)
-		hashKeys := generateHashKeys(rowNum)
+		hashKeys := testutils.GenerateHashKeys(rowNum)
 		return &milvuspb.UpsertRequest{
 			Base:           nil,
 			DbName:         dbName,
@@ -3755,7 +3756,7 @@ func TestProxy(t *testing.T) {
 		pkFieldData := newScalarFieldData(schema.Fields[0], int64Field, rowNum)
 		fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
 		bVecColumn := newBinaryVectorFieldData(binaryVecField, rowNum, dim)
-		hashKeys := generateHashKeys(rowNum)
+		hashKeys := testutils.GenerateHashKeys(rowNum)
 		return &milvuspb.UpsertRequest{
 			Base:           nil,
 			DbName:         dbName,

diff --git a/internal/proxy/task_insert_test.go b/internal/proxy/task_insert_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/mq/msgstream"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/testutils"
 )
 
 func TestInsertTask_CheckAligned(t *testing.T) {
@@ -56,8 +57,8 @@ func TestInsertTask_CheckAligned(t *testing.T) {
 					MsgType: commonpb.MsgType_Insert,
 				},
 				Version:    msgpb.InsertDataVersion_ColumnBased,
-				RowIDs:     generateInt64Array(numRows),
-				Timestamps: generateUint64Array(numRows),
+				RowIDs:     testutils.GenerateInt64Array(numRows),
+				Timestamps: testutils.GenerateUint64Array(numRows),
 			},
 		},
 		schema: &schemapb.CollectionSchema{

diff --git a/internal/proxy/task_query_test.go b/internal/proxy/task_query_test.go
@@ -39,6 +39,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/funcutil"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/testutils"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
@@ -195,7 +196,7 @@ func TestQueryTask_all(t *testing.T) {
 		Status: merr.Success(),
 		Ids: &schemapb.IDs{
 			IdField: &schemapb.IDs_IntId{
-				IntId: &schemapb.LongArray{Data: generateInt64Array(hitNum)},
+				IntId: &schemapb.LongArray{Data: testutils.GenerateInt64Array(hitNum)},
 			},
 		},
 	}

diff --git a/internal/proxy/task_test.go b/internal/proxy/task_test.go
@@ -49,6 +49,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/metric"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/testutils"
 	"github.com/milvus-io/milvus/pkg/util/timerecord"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 	"github.com/milvus-io/milvus/pkg/util/uniquegenerator"
@@ -1680,7 +1681,7 @@ func TestTask_Int64PrimaryKey(t *testing.T) {
 	defer segAllocator.Close()
 
 	t.Run("insert", func(t *testing.T) {
-		hash := generateHashKeys(nb)
+		hash := testutils.GenerateHashKeys(nb)
 		task := &insertTask{
 			insertMsg: &BaseInsertTask{
 				BaseMsg: msgstream.BaseMsg{
@@ -1874,7 +1875,7 @@ func TestTask_VarCharPrimaryKey(t *testing.T) {
 	defer segAllocator.Close()
 
 	t.Run("insert", func(t *testing.T) {
-		hash := generateHashKeys(nb)
+		hash := testutils.GenerateHashKeys(nb)
 		task := &insertTask{
 			insertMsg: &BaseInsertTask{
 				BaseMsg: msgstream.BaseMsg{
@@ -1929,7 +1930,7 @@ func TestTask_VarCharPrimaryKey(t *testing.T) {
 	})
 
 	t.Run("upsert", func(t *testing.T) {
-		hash := generateHashKeys(nb)
+		hash := testutils.GenerateHashKeys(nb)
 		task := &upsertTask{
 			upsertMsg: &msgstream.UpsertMsg{
 				InsertMsg: &BaseInsertTask{
@@ -3339,7 +3340,7 @@ func TestPartitionKey(t *testing.T) {
 	})
 
 	t.Run("Upsert", func(t *testing.T) {
-		hash := generateHashKeys(nb)
+		hash := testutils.GenerateHashKeys(nb)
 		ut := &upsertTask{
 			ctx:       ctx,
 			Condition: NewTaskCondition(ctx),

diff --git a/internal/proxy/task_upsert_test.go b/internal/proxy/task_upsert_test.go
@@ -28,6 +28,7 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
 	"github.com/milvus-io/milvus/pkg/mq/msgstream"
 	"github.com/milvus-io/milvus/pkg/util/commonpbutil"
+	"github.com/milvus-io/milvus/pkg/util/testutils"
 )
 
 func TestUpsertTask_CheckAligned(t *testing.T) {
@@ -96,8 +97,8 @@ func TestUpsertTask_CheckAligned(t *testing.T) {
 			NumRows:    uint32(numRows),
 			FieldsData: []*schemapb.FieldData{},
 		},
-		rowIDs:     generateInt64Array(numRows),
-		timestamps: generateUint64Array(numRows),
+		rowIDs:     testutils.GenerateInt64Array(numRows),
+		timestamps: testutils.GenerateUint64Array(numRows),
 		schema:     schema,
 		upsertMsg: &msgstream.UpsertMsg{
 			InsertMsg: &msgstream.InsertMsg{

diff --git a/internal/proxy/validate_util_test.go b/internal/proxy/validate_util_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
 	"github.com/milvus-io/milvus/pkg/common"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/testutils"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
@@ -291,8 +292,8 @@ func Test_validateUtil_checkFloatVectorFieldData(t *testing.T) {
 func Test_validateUtil_checkFloat16VectorFieldData(t *testing.T) {
 	nb := 5
 	dim := int64(8)
-	data := generateFloat16Vectors(nb, int(dim))
-	invalidData := generateFloat16VectorsWithInvalidData(nb, int(dim))
+	data := testutils.GenerateFloat16Vectors(nb, int(dim))
+	invalidData := testutils.GenerateFloat16VectorsWithInvalidData(nb, int(dim))
 
 	t.Run("not float16 vector", func(t *testing.T) {
 		f := &schemapb.FieldData{}
@@ -382,8 +383,8 @@ func Test_validateUtil_checkFloat16VectorFieldData(t *testing.T) {
 func Test_validateUtil_checkBfloatVectorFieldData(t *testing.T) {
 	nb := 5
 	dim := int64(8)
-	data := generateFloat16Vectors(nb, int(dim))
-	invalidData := generateBFloat16VectorsWithInvalidData(nb, int(dim))
+	data := testutils.GenerateFloat16Vectors(nb, int(dim))
+	invalidData := testutils.GenerateBFloat16VectorsWithInvalidData(nb, int(dim))
 	t.Run("not float vector", func(t *testing.T) {
 		f := &schemapb.FieldData{}
 		v := newValidateUtil()
@@ -1196,7 +1197,7 @@ func Test_validateUtil_checkAligned(t *testing.T) {
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_FloatVector{
 							FloatVector: &schemapb.FloatArray{
-								Data: generateFloatVectors(10, 8),
+								Data: testutils.GenerateFloatVectors(10, 8),
 							},
 						},
 						Dim: 8,
@@ -1209,7 +1210,7 @@ func Test_validateUtil_checkAligned(t *testing.T) {
 				Field: &schemapb.FieldData_Vectors{
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_BinaryVector{
-							BinaryVector: generateBinaryVectors(10, 8),
+							BinaryVector: testutils.GenerateBinaryVectors(10, 8),
 						},
 						Dim: 8,
 					},
@@ -1222,7 +1223,7 @@ func Test_validateUtil_checkAligned(t *testing.T) {
 					Scalars: &schemapb.ScalarField{
 						Data: &schemapb.ScalarField_StringData{
 							StringData: &schemapb.StringArray{
-								Data: generateVarCharArray(10, 8),
+								Data: testutils.GenerateVarCharArray(10, 8),
 							},
 						},
 					},
@@ -1356,7 +1357,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 				Field: &schemapb.FieldData_Vectors{
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_BinaryVector{
-							BinaryVector: generateBinaryVectors(2, 8),
+							BinaryVector: testutils.GenerateBinaryVectors(2, 8),
 						},
 					},
 				},
@@ -1368,7 +1369,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 					Scalars: &schemapb.ScalarField{
 						Data: &schemapb.ScalarField_StringData{
 							StringData: &schemapb.StringArray{
-								Data: generateVarCharArray(2, 8),
+								Data: testutils.GenerateVarCharArray(2, 8),
 							},
 						},
 					},
@@ -1380,7 +1381,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 				Field: &schemapb.FieldData_Vectors{
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_Float16Vector{
-							Float16Vector: generateFloat16Vectors(2, 8),
+							Float16Vector: testutils.GenerateFloat16Vectors(2, 8),
 						},
 					},
 				},
@@ -1391,7 +1392,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 				Field: &schemapb.FieldData_Vectors{
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_Bfloat16Vector{
-							Bfloat16Vector: generateBFloat16Vectors(2, 8),
+							Bfloat16Vector: testutils.GenerateBFloat16Vectors(2, 8),
 						},
 					},
 				},
@@ -1474,7 +1475,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_FloatVector{
 							FloatVector: &schemapb.FloatArray{
-								Data: generateFloatVectors(2, 1),
+								Data: testutils.GenerateFloatVectors(2, 1),
 							},
 						},
 					},
@@ -1486,7 +1487,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 				Field: &schemapb.FieldData_Vectors{
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_BinaryVector{
-							BinaryVector: generateBinaryVectors(2, 8),
+							BinaryVector: testutils.GenerateBinaryVectors(2, 8),
 						},
 					},
 				},
@@ -1497,7 +1498,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 				Field: &schemapb.FieldData_Vectors{
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_Float16Vector{
-							Float16Vector: generateFloat16Vectors(2, 8),
+							Float16Vector: testutils.GenerateFloat16Vectors(2, 8),
 						},
 					},
 				},
@@ -1508,7 +1509,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 				Field: &schemapb.FieldData_Vectors{
 					Vectors: &schemapb.VectorField{
 						Data: &schemapb.VectorField_Bfloat16Vector{
-							Bfloat16Vector: generateBFloat16Vectors(2, 8),
+							Bfloat16Vector: testutils.GenerateBFloat16Vectors(2, 8),
 						},
 					},
 				},
@@ -2323,7 +2324,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 						Dim: 8,
 						Data: &schemapb.VectorField_FloatVector{
 							FloatVector: &schemapb.FloatArray{
-								Data: generateFloatVectors(2, 8),
+								Data: testutils.GenerateFloatVectors(2, 8),
 							},
 						},
 					},
@@ -2336,7 +2337,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 					Vectors: &schemapb.VectorField{
 						Dim: 8,
 						Data: &schemapb.VectorField_BinaryVector{
-							BinaryVector: generateBinaryVectors(2, 8),
+							BinaryVector: testutils.GenerateBinaryVectors(2, 8),
 						},
 					},
 				},
@@ -2348,7 +2349,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 					Scalars: &schemapb.ScalarField{
 						Data: &schemapb.ScalarField_StringData{
 							StringData: &schemapb.StringArray{
-								Data: generateVarCharArray(2, 8),
+								Data: testutils.GenerateVarCharArray(2, 8),
 							},
 						},
 					},
@@ -2568,7 +2569,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 					Scalars: &schemapb.ScalarField{
 						Data: &schemapb.ScalarField_FloatData{
 							FloatData: &schemapb.FloatArray{
-								Data: generateFloat32Array(2),
+								Data: testutils.GenerateFloat32Array(2),
 							},
 						},
 					},
@@ -2581,7 +2582,7 @@ func Test_validateUtil_Validate(t *testing.T) {
 					Scalars: &schemapb.ScalarField{
 						Data: &schemapb.ScalarField_DoubleData{
 							DoubleData: &schemapb.DoubleArray{
-								Data: generateFloat64Array(2),
+								Data: testutils.GenerateFloat64Array(2),
 							},
 						},
 					},

diff --git a/internal/querynodev2/pipeline/mock_data.go b/internal/querynodev2/pipeline/mock_data.go
@@ -22,9 +22,9 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
-	"github.com/milvus-io/milvus/internal/querynodev2/segments"
 	"github.com/milvus-io/milvus/pkg/mq/msgstream"
 	"github.com/milvus-io/milvus/pkg/util/commonpbutil"
+	"github.com/milvus-io/milvus/pkg/util/testutils"
 )
 
 const defaultDim = 128
@@ -164,9 +164,9 @@ func genFiledDataWithSchema(schema *schemapb.CollectionSchema, numRows int) []*s
 	fieldsData := make([]*schemapb.FieldData, 0)
 	for _, field := range schema.Fields {
 		if field.DataType < 100 {
-			fieldsData = append(fieldsData, segments.GenTestScalarFieldData(field.DataType, field.DataType.String(), field.GetFieldID(), numRows))
+			fieldsData = append(fieldsData, testutils.GenerateScalarFieldDataWithID(field.DataType, field.DataType.String(), field.GetFieldID(), numRows))
 		} else {
-			fieldsData = append(fieldsData, segments.GenTestVectorFiledData(field.DataType, field.DataType.String(), field.GetFieldID(), numRows, defaultDim))
+			fieldsData = append(fieldsData, testutils.GenerateVectorFieldDataWithID(field.DataType, field.DataType.String(), field.GetFieldID(), numRows, defaultDim))
 		}
 	}
 	return fieldsData