From 5684abd27e4f51512362116042c81be60194417b Mon Sep 17 00:00:00 2001
From: Bin Shi <binshi.bing@gmail.com>
Date: Mon, 3 Apr 2023 19:19:52 -0700
Subject: [PATCH] Simplify test

Signed-off-by: Bin Shi <binshi.bing@gmail.com>
---
 pkg/tso/keyspace_group_manager_test.go | 74 ++++++++------------------
 1 file changed, 21 insertions(+), 53 deletions(-)

diff --git a/pkg/tso/keyspace_group_manager_test.go b/pkg/tso/keyspace_group_manager_test.go
index d1215a59fae4..989ddb0b40fd 100644
--- a/pkg/tso/keyspace_group_manager_test.go
+++ b/pkg/tso/keyspace_group_manager_test.go
@@ -17,8 +17,10 @@ package tso
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"math/rand"
 	"path"
+	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -174,13 +176,13 @@ func runTestLoadKeyspaceGroupsAssignment(
 	loadKeyspaceGroupsBatchSize int64, // set to 0 to use the default value
 	probabilityAssignToMe int, // percentage of assigning keyspace groups to this host/pod
 ) {
-	ids := make(map[uint32]bool, 0)
+	idsExpected := []int{}
 	mgr := newUniqueKeyspaceGroupManager(ctx, etcdClient, cfg, 0, loadKeyspaceGroupsBatchSize)
 	re.NotNil(mgr)
 	defer mgr.Close()
 
-	const step = 30
-	keyspaceGroupsAdded := sync.Map{}
+	step := 30
+	mux := sync.Mutex{}
 	wg := sync.WaitGroup{}
 	for i := 0; i < numberOfKeypaceGroupsToAdd; i += step {
 		wg.Add(1)
@@ -197,7 +199,9 @@ func runTestLoadKeyspaceGroupsAssignment(
 				assignToMe := false
 				if randomGen.Intn(100) < probabilityAssignToMe {
 					assignToMe = true
-					keyspaceGroupsAdded.Store(uint32(j), struct{}{})
+					mux.Lock()
+					idsExpected = append(idsExpected, j)
+					mux.Unlock()
 				}
 				addKeyspaceGroupAssignment(
 					ctx, etcdClient, assignToMe,
@@ -207,21 +211,13 @@ func runTestLoadKeyspaceGroupsAssignment(
 	}
 	wg.Wait()
 
-	// Set the expected result. The keyspace group manager only loads len(mgr.ams) keyspace
-	// groups as max, so we need to ignore the extra keyspace groups in the expected result.
-	keyspaceGroupsAdded.Range(func(key, _ interface{}) bool {
-		id := key.(uint32)
-		if id < uint32(numberOfKeypaceGroupsToAdd) {
-			ids[id] = true
-		}
-		return true
-	})
-
 	err := mgr.Initialize(true)
 	re.NoError(err)
 
 	// Verify the keyspace group assignment.
-	re.True(verifyKeyspaceGroupAssignment(ids, mgr))
+	sort.Ints(idsExpected)
+	idsAssigned := collectAssignedKeyspaceGroupIDs(re, mgr)
+	re.Equal(idsExpected, idsAssigned)
 }
 
 func newUniqueKeyspaceGroupManager(
@@ -277,50 +273,22 @@ func addKeyspaceGroupAssignment(
 	return nil
 }
 
-func verifyKeyspaceGroupAssignment(ids map[uint32]bool, ksgMgr *KeyspaceGroupManager) bool {
-	for i := 0; i < len(ksgMgr.ams); i++ {
-		am := ksgMgr.ams[i].Load()
-		if am == nil {
-			if _, ok := ids[uint32(i)]; ok {
-				return false
-			}
-		} else {
-			if i != int(am.ksgID) {
-				return false
-			}
-			if _, ok := ids[uint32(i)]; !ok {
-				return false
-			}
-		}
-	}
-
+func collectAssignedKeyspaceGroupIDs(re *require.Assertions, ksgMgr *KeyspaceGroupManager) []int {
+	ids := []int{}
 	for i := 0; i < len(ksgMgr.ksgs); i++ {
 		ksg := ksgMgr.ksgs[i].Load()
 		if ksg == nil {
-			if _, ok := ids[uint32(i)]; ok {
-				return false
-			}
+			re.Nil(ksgMgr.ams[i].Load(), fmt.Sprintf("ksg is nil but am is not nil for id %d", i))
 		} else {
-			if i != int(ksg.ID) {
-				return false
-			}
-			if _, ok := ids[uint32(i)]; !ok {
-				return false
-			}
-			if ksg.Members[0].Address != ksgMgr.tsoServiceID.ServiceAddr {
-				return false
+			am := ksgMgr.ams[i].Load()
+			re.NotNil(am, fmt.Sprintf("ksg is not nil but am is nil for id %d", i))
+			re.Equal(i, int(am.ksgID))
+			re.Equal(i, int(ksg.ID))
+			if ksg.Members[0].Address == ksgMgr.tsoServiceID.ServiceAddr {
+				ids = append(ids, i)
 			}
-			if len(ksg.Keyspaces) != 1 || ksg.Keyspaces[0] != uint32(i) {
-				return false
-			}
-		}
-	}
-
-	for k := range ids {
-		if k >= uint32(len(ksgMgr.ams)) || ksgMgr.ams[k].Load() == nil || ksgMgr.ksgs[k].Load() == nil {
-			return false
 		}
 	}
 
-	return true
+	return ids
 }