From fdd9561726757c8270fbc355c0eb08421c70807b Mon Sep 17 00:00:00 2001
From: Lulu Zhang <lulu@buildbuddy.io>
Date: Wed, 18 Sep 2024 12:29:02 -0700
Subject: [PATCH] [raft] add a test for ConfiguredClusters in the store (#7488)

---
 enterprise/server/raft/store/store.go       |  1 +
 enterprise/server/raft/store/store_test.go  | 10 ++++++
 enterprise/server/raft/testutil/testutil.go | 34 +++++++++++----------
 3 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/enterprise/server/raft/store/store.go b/enterprise/server/raft/store/store.go
index b7ed9f439dd..f6fd9ef0269 100644
--- a/enterprise/server/raft/store/store.go
+++ b/enterprise/server/raft/store/store.go
@@ -278,6 +278,7 @@ func NewWithArgs(env environment.Env, rootDir string, nodeHost *dragonboat.NodeH
 			if err := nodeHost.StartOnDiskReplica(nil, false /*=join*/, s.ReplicaFactoryFn, rc); err != nil {
 				return nil, status.InternalErrorf("failed to start c%dn%d: %s", logInfo.ShardID, logInfo.ReplicaID, err)
 			}
+			s.configuredClusters++
 		} else {
 			replicaDescriptor := &rfpb.ReplicaDescriptor{RangeId: logInfo.ShardID, ReplicaId: logInfo.ReplicaID}
 			previouslyStartedReplicas = append(previouslyStartedReplicas, replicaDescriptor)
diff --git a/enterprise/server/raft/store/store_test.go b/enterprise/server/raft/store/store_test.go
index eacc2af0410..d96a8e42f9d 100644
--- a/enterprise/server/raft/store/store_test.go
+++ b/enterprise/server/raft/store/store_test.go
@@ -56,6 +56,16 @@ func getMembership(t *testing.T, ts *testutil.TestingStore, ctx context.Context,
 	return replicas
 }
 
+func TestConfiguredClusters(t *testing.T) {
+	sf := testutil.NewStoreFactory(t)
+	s1 := sf.NewStore(t)
+	ctx := context.Background()
+	sf.StartShard(t, ctx, s1)
+	s1.Stop()
+	sf.RecreateStore(t, s1)
+	require.Equal(t, 2, s1.ConfiguredClusters())
+}
+
 func TestAddGetRemoveRange(t *testing.T) {
 	sf := testutil.NewStoreFactory(t)
 	s1 := sf.NewStore(t)
diff --git a/enterprise/server/raft/testutil/testutil.go b/enterprise/server/raft/testutil/testutil.go
index 93b3350a28b..69a5c28e033 100644
--- a/enterprise/server/raft/testutil/testutil.go
+++ b/enterprise/server/raft/testutil/testutil.go
@@ -41,7 +41,6 @@ func localAddr(t *testing.T) string {
 
 type StoreFactory struct {
 	rootDir     string
-	fileDir     string
 	gossipAddrs []string
 	reg         registry.NodeRegistry
 	clock       clockwork.Clock
@@ -58,7 +57,6 @@ func NewStoreFactoryWithClock(t *testing.T, clock clockwork.Clock) *StoreFactory
 	require.NoError(t, err)
 	return &StoreFactory{
 		rootDir: rootDir,
-		fileDir: fileDir,
 		reg:     registry.NewStaticNodeRegistry(1, nil),
 		clock:   clock,
 	}
@@ -74,19 +72,7 @@ func (sf *StoreFactory) Registry() registry.NodeRegistry {
 	return sf.reg
 }
 
-func (sf *StoreFactory) NewStore(t *testing.T) *TestingStore {
-	nodeAddr := localAddr(t)
-	gm, err := gossip.New("name-"+nodeAddr, nodeAddr, sf.gossipAddrs)
-	require.NoError(t, err)
-	sf.gossipAddrs = append(sf.gossipAddrs, nodeAddr)
-
-	ts := &TestingStore{
-		t:           t,
-		gm:          gm,
-		RaftAddress: localAddr(t),
-		GRPCAddress: localAddr(t),
-		RootDir:     filepath.Join(sf.rootDir, fmt.Sprintf("store-%d", len(sf.gossipAddrs))),
-	}
+func (sf *StoreFactory) RecreateStore(t *testing.T, ts *TestingStore) {
 	require.Nil(t, disk.EnsureDirectoryExists(ts.RootDir))
 
 	reg := sf.reg
@@ -127,7 +113,7 @@ func (sf *StoreFactory) NewStore(t *testing.T) *TestingStore {
 	require.NoError(t, err)
 	leaser := pebble.NewDBLeaser(db)
 	ts.leaser = leaser
-	store, err := store.NewWithArgs(te, ts.RootDir, nodeHost, gm, s, reg, raftListener, apiClient, ts.GRPCAddress, partitions, db, leaser)
+	store, err := store.NewWithArgs(te, ts.RootDir, nodeHost, ts.gm, s, reg, raftListener, apiClient, ts.GRPCAddress, partitions, db, leaser)
 	require.NoError(t, err)
 	require.NotNil(t, store)
 	store.Start()
@@ -136,6 +122,22 @@ func (sf *StoreFactory) NewStore(t *testing.T) *TestingStore {
 	t.Cleanup(func() {
 		ts.Stop()
 	})
+}
+
+func (sf *StoreFactory) NewStore(t *testing.T) *TestingStore {
+	nodeAddr := localAddr(t)
+	gm, err := gossip.New("name-"+nodeAddr, nodeAddr, sf.gossipAddrs)
+	require.NoError(t, err)
+	sf.gossipAddrs = append(sf.gossipAddrs, nodeAddr)
+
+	ts := &TestingStore{
+		t:           t,
+		gm:          gm,
+		RaftAddress: localAddr(t),
+		GRPCAddress: localAddr(t),
+		RootDir:     filepath.Join(sf.rootDir, fmt.Sprintf("store-%d", len(sf.gossipAddrs))),
+	}
+	sf.RecreateStore(t, ts)
 	return ts
 }