diff --git a/internal/master/runtime_stats_test.go b/internal/master/runtime_stats_test.go
index 539bd2894be71069a1d5742cf5d2eac5edb15bb1..7e9949cec10d60dc70566c9e0ce01fa7dc36d60f 100644
--- a/internal/master/runtime_stats_test.go
+++ b/internal/master/runtime_stats_test.go
@@ -41,3 +41,33 @@ func TestRuntimeStats_UpdateFieldStats(t *testing.T) {
 		assert.EqualValues(t, 1, found)
 	}
 }
+
+func TestRuntimeStats_GetTotalNumOfRelatedSegments(t *testing.T) {
+	runtimeStats := NewRuntimeStats()
+	runtimeStats.collStats = make(map[UniqueID]*CollRuntimeStats)
+
+	runtimeStats.collStats[1] = &CollRuntimeStats{
+		fieldIndexStats: map[UniqueID][]*FieldIndexRuntimeStats{
+			100: {
+				{1, []*commonpb.KeyValuePair{{Key: "k1", Value: "v1"}}, 10},
+				{3, []*commonpb.KeyValuePair{{Key: "k1", Value: "v1"}}, 20},
+				{2, []*commonpb.KeyValuePair{{Key: "k2", Value: "v2"}}, 20},
+			},
+			200: {
+				{1, []*commonpb.KeyValuePair{}, 20},
+			},
+		},
+	}
+
+	runtimeStats.collStats[2] = &CollRuntimeStats{
+		fieldIndexStats: map[UniqueID][]*FieldIndexRuntimeStats{
+			100: {
+				{1, []*commonpb.KeyValuePair{{Key: "k1", Value: "v1"}}, 10},
+			},
+		},
+	}
+	assert.EqualValues(t, 30, runtimeStats.GetTotalNumOfRelatedSegments(1, 100, []*commonpb.KeyValuePair{{Key: "k1", Value: "v1"}}))
+	assert.EqualValues(t, 20, runtimeStats.GetTotalNumOfRelatedSegments(1, 100, []*commonpb.KeyValuePair{{Key: "k2", Value: "v2"}}))
+	assert.EqualValues(t, 20, runtimeStats.GetTotalNumOfRelatedSegments(1, 200, []*commonpb.KeyValuePair{}))
+	assert.EqualValues(t, 10, runtimeStats.GetTotalNumOfRelatedSegments(2, 100, []*commonpb.KeyValuePair{{Key: "k1", Value: "v1"}}))
+}