Skip to content

Commit 389e986

Browse files
committed
add test
Signed-off-by: aoiasd <[email protected]>
1 parent 0358144 commit 389e986

File tree

2 files changed

+13
-1
lines changed

2 files changed

+13
-1
lines changed

internal/flushcommon/pipeline/flow_graph_embedding_node_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
7373
t.Run("normal case", func(t *testing.T) {
7474
node, err := newEmbeddingNode("test-channel", metaCache)
7575
assert.NoError(t, err)
76+
defer node.Close()
7677

7778
var output []Msg
7879
assert.NotPanics(t, func() {
@@ -114,6 +115,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
114115
t.Run("with close msg", func(t *testing.T) {
115116
node, err := newEmbeddingNode("test-channel", metaCache)
116117
assert.NoError(t, err)
118+
defer node.Close()
117119

118120
var output []Msg
119121

@@ -131,6 +133,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
131133
t.Run("prepare insert failed", func(t *testing.T) {
132134
node, err := newEmbeddingNode("test-channel", metaCache)
133135
assert.NoError(t, err)
136+
defer node.Close()
134137

135138
assert.Panics(t, func() {
136139
node.Operate([]Msg{
@@ -152,6 +155,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
152155
t.Run("embedding failed", func(t *testing.T) {
153156
node, err := newEmbeddingNode("test-channel", metaCache)
154157
assert.NoError(t, err)
158+
defer node.Close()
155159

156160
node.functionRunners[0].GetSchema().Type = 0
157161
assert.Panics(t, func() {

internal/querynodev2/pipeline/embedding_node_test.go

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,8 +140,9 @@ func (suite *EmbeddingNodeSuite) TestCreateEmbeddingNode() {
140140
collSchema := proto.Clone(suite.collectionSchema).(*schemapb.CollectionSchema)
141141
collection := segments.NewCollectionWithoutSegcoreForTest(suite.collectionID, collSchema)
142142
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
143-
_, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
143+
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
144144
suite.NoError(err)
145+
defer node.Close()
145146
})
146147
}
147148

@@ -151,6 +152,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
151152
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
152153
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
153154
suite.NoError(err)
155+
defer node.Close()
154156

155157
suite.colManager.EXPECT().Get(suite.collectionID).Return(nil).Once()
156158
suite.Panics(func() {
@@ -163,6 +165,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
163165
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Times(2)
164166
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
165167
suite.NoError(err)
168+
defer node.Close()
166169

167170
suite.Panics(func() {
168171
node.Operate(&insertNodeMsg{
@@ -192,6 +195,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
192195
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Times(2)
193196
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
194197
suite.NoError(err)
198+
defer node.Close()
195199

196200
suite.NotPanics(func() {
197201
output := node.Operate(&insertNodeMsg{
@@ -213,6 +217,7 @@ func (suite *EmbeddingNodeSuite) TestAddInsertData() {
213217
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
214218
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
215219
suite.NoError(err)
220+
defer node.Close()
216221

217222
// transfer insert msg failed because rowbase data not support sparse vector
218223
insertDatas := make(map[int64]*delegator.InsertData)
@@ -237,6 +242,7 @@ func (suite *EmbeddingNodeSuite) TestAddInsertData() {
237242
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
238243
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
239244
suite.NoError(err)
245+
defer node.Close()
240246

241247
insertDatas := make(map[int64]*delegator.InsertData)
242248
err = node.addInsertData(insertDatas, suite.msgs[0], collection)
@@ -250,6 +256,7 @@ func (suite *EmbeddingNodeSuite) TestBM25Embedding() {
250256
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
251257
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
252258
suite.NoError(err)
259+
defer node.Close()
253260

254261
runner := function.NewMockFunctionRunner(suite.T())
255262
runner.EXPECT().BatchRun(mock.Anything).Return(nil, errors.New("mock error"))
@@ -265,6 +272,7 @@ func (suite *EmbeddingNodeSuite) TestBM25Embedding() {
265272
suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
266273
node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
267274
suite.NoError(err)
275+
defer node.Close()
268276

269277
runner := function.NewMockFunctionRunner(suite.T())
270278
runner.EXPECT().BatchRun(mock.Anything).Return([]interface{}{1}, nil)

0 commit comments

Comments
 (0)