From 31c5b286da5226e520e4cb15a76b4a7e062ac394 Mon Sep 17 00:00:00 2001
From: yah01 <yang.cen@zilliz.com>
Date: Fri, 18 Mar 2022 13:17:22 +0800
Subject: [PATCH] Fix wrong unittest, check whether loading collection will
 cause OOM (#16068)

Signed-off-by: yah01 <yang.cen@zilliz.com>
---
 internal/querynode/task_test.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/internal/querynode/task_test.go b/internal/querynode/task_test.go
index b77dd5826c..8a9a8bf868 100644
--- a/internal/querynode/task_test.go
+++ b/internal/querynode/task_test.go
@@ -634,10 +634,16 @@ func TestTask_loadSegmentsTask(t *testing.T) {
 				PartitionID:  defaultPartitionID,
 				CollectionID: defaultCollectionID,
 				NumOfRows:    totalRAM / int64(sizePerRecord),
+				SegmentSize:  totalRAM,
 			},
 		}
+		// Reach the segment size that would cause OOM
+		for node.loader.checkSegmentSize(defaultCollectionID, task.req.Infos, 1) == nil {
+			task.req.Infos[0].SegmentSize *= 2
+		}
 		err = task.Execute(ctx)
 		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "OOM")
 	})
 }