add test

Signed-off-by: aoiasd <[email protected]>

Author: aoiasd

internal/flushcommon/pipeline/flow_graph_embedding_node_test.go

@@ -73,6 +73,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
 	t.Run("normal case", func(t *testing.T) {
 		node, err := newEmbeddingNode("test-channel", metaCache)
 		assert.NoError(t, err)
+		defer node.Close()
 
 		var output []Msg
 		assert.NotPanics(t, func() {
@@ -114,6 +115,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
 	t.Run("with close msg", func(t *testing.T) {
 		node, err := newEmbeddingNode("test-channel", metaCache)
 		assert.NoError(t, err)
+		defer node.Close()
 
 		var output []Msg
 
@@ -131,6 +133,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
 	t.Run("prepare insert failed", func(t *testing.T) {
 		node, err := newEmbeddingNode("test-channel", metaCache)
 		assert.NoError(t, err)
+		defer node.Close()
 
 		assert.Panics(t, func() {
 			node.Operate([]Msg{
@@ -152,6 +155,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
 	t.Run("embedding failed", func(t *testing.T) {
 		node, err := newEmbeddingNode("test-channel", metaCache)
 		assert.NoError(t, err)
+		defer node.Close()
 
 		node.functionRunners[0].GetSchema().Type = 0
 		assert.Panics(t, func() {

internal/querynodev2/pipeline/embedding_node_test.go

@@ -140,8 +140,9 @@ func (suite *EmbeddingNodeSuite) TestCreateEmbeddingNode() {
 		collSchema := proto.Clone(suite.collectionSchema).(*schemapb.CollectionSchema)
 		collection := segments.NewCollectionWithoutSegcoreForTest(suite.collectionID, collSchema)
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
-		_, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
+		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 	})
 }
 
@@ -151,6 +152,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(nil).Once()
 		suite.Panics(func() {
@@ -163,6 +165,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Times(2)
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		suite.Panics(func() {
 			node.Operate(&insertNodeMsg{
@@ -192,6 +195,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Times(2)
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		suite.NotPanics(func() {
 			output := node.Operate(&insertNodeMsg{
@@ -213,6 +217,7 @@ func (suite *EmbeddingNodeSuite) TestAddInsertData() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		// transfer insert msg failed because rowbase data not support sparse vector
 		insertDatas := make(map[int64]*delegator.InsertData)
@@ -237,6 +242,7 @@ func (suite *EmbeddingNodeSuite) TestAddInsertData() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		insertDatas := make(map[int64]*delegator.InsertData)
 		err = node.addInsertData(insertDatas, suite.msgs[0], collection)
@@ -250,6 +256,7 @@ func (suite *EmbeddingNodeSuite) TestBM25Embedding() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		runner := function.NewMockFunctionRunner(suite.T())
 		runner.EXPECT().BatchRun(mock.Anything).Return(nil, errors.New("mock error"))
@@ -265,6 +272,7 @@ func (suite *EmbeddingNodeSuite) TestBM25Embedding() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		runner := function.NewMockFunctionRunner(suite.T())
 		runner.EXPECT().BatchRun(mock.Anything).Return([]interface{}{1}, nil)