*: Add unit test for sort hot peers (#9276)

ref #4399

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

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

Author: okJiang

pkg/schedule/schedulers/hot_region_solver.go

@@ -437,7 +437,7 @@ func (bs *balanceSolver) filterHotPeers(storeLoad *statistics.StoreLoadDetail) [
 		bs.nthHotPeer[storeID][bs.secondPriority] = secondSort[topnPosition-1]
 	}
 	if len(hotPeers) > bs.maxPeerNum {
-		union := bs.sortHotPeers(firstSort, secondSort)
+		union := sortHotPeers(firstSort, secondSort, bs.maxPeerNum)
 		ret = make([]*statistics.HotPeerStat, 0, len(union))
 		for peer := range union {
 			appendItem(peer)
@@ -451,10 +451,13 @@ func (bs *balanceSolver) filterHotPeers(storeLoad *statistics.StoreLoadDetail) [
 	return ret
 }
 
-func (bs *balanceSolver) sortHotPeers(firstSort, secondSort []*statistics.HotPeerStat) map[*statistics.HotPeerStat]struct{} {
-	union := make(map[*statistics.HotPeerStat]struct{}, bs.maxPeerNum)
+func sortHotPeers[T any](firstSort, secondSort []*T, maxPeerNum int) map[*T]struct{} {
+	union := make(map[*T]struct{}, maxPeerNum)
 	// At most MaxPeerNum peers, to prevent balanceSolver.solve() too slow.
-	for len(union) < bs.maxPeerNum {
+	for len(union) < maxPeerNum {
+		if len(firstSort) == 0 && len(secondSort) == 0 {
+			break
+		}
 		for len(firstSort) > 0 {
 			peer := firstSort[0]
 			firstSort = firstSort[1:]
@@ -463,7 +466,7 @@ func (bs *balanceSolver) sortHotPeers(firstSort, secondSort []*statistics.HotPee
 				break
 			}
 		}
-		for len(union) < bs.maxPeerNum && len(secondSort) > 0 {
+		for len(union) < maxPeerNum && len(secondSort) > 0 {
 			peer := secondSort[0]
 			secondSort = secondSort[1:]
 			if _, ok := union[peer]; !ok {

pkg/schedule/schedulers/hot_region_solver_test.go

@@ -606,3 +606,79 @@ func TestBucketFirstStat(t *testing.T) {
 		re.Equal(data.expect, bs.bucketFirstStat())
 	}
 }
+
+func TestSortHotPeers(t *testing.T) {
+	re := require.New(t)
+
+	type testPeer struct {
+		id int
+	}
+
+	peer1 := &testPeer{id: 1}
+	peer2 := &testPeer{id: 2}
+	peer3 := &testPeer{id: 3}
+	peer4 := &testPeer{id: 4}
+
+	tests := []struct {
+		name       string
+		firstSort  []*testPeer
+		secondSort []*testPeer
+		maxPeerNum int
+		expected   []*testPeer
+	}{
+		{
+			name:       "No duplicates, maxPeerNum greater than total peers",
+			firstSort:  []*testPeer{peer1, peer2},
+			secondSort: []*testPeer{peer3, peer4},
+			maxPeerNum: 5,
+			expected:   []*testPeer{peer1, peer3, peer2, peer4},
+		},
+		{
+			name:       "No duplicates, maxPeerNum less than total peers",
+			firstSort:  []*testPeer{peer1, peer2},
+			secondSort: []*testPeer{peer3, peer4},
+			maxPeerNum: 3,
+			expected:   []*testPeer{peer1, peer3, peer2},
+		},
+		{
+			name:       "Duplicates in both lists",
+			firstSort:  []*testPeer{peer1, peer2},
+			secondSort: []*testPeer{peer2, peer3},
+			maxPeerNum: 3,
+			expected:   []*testPeer{peer1, peer2, peer3},
+		},
+		{
+			name:       "Empty firstSort",
+			firstSort:  []*testPeer{},
+			secondSort: []*testPeer{peer3, peer4},
+			maxPeerNum: 2,
+			expected:   []*testPeer{peer3, peer4},
+		},
+		{
+			name:       "Empty secondSort",
+			firstSort:  []*testPeer{peer1, peer2},
+			secondSort: []*testPeer{},
+			maxPeerNum: 2,
+			expected:   []*testPeer{peer1, peer2},
+		},
+		{
+			name:       "Both lists empty",
+			firstSort:  []*testPeer{},
+			secondSort: []*testPeer{},
+			maxPeerNum: 2,
+			expected:   []*testPeer{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(*testing.T) {
+			result := sortHotPeers(tt.firstSort, tt.secondSort, tt.maxPeerNum)
+			re.Len(result, len(tt.expected))
+
+			for _, expectedPeer := range tt.expected {
+				_, exists := result[expectedPeer]
+				re.True(exists, "Expected peer not found in result")
+			}
+		})
+	}
+}