From e9fa3d30d714c9251a01a68d373c476d5bbf31ab Mon Sep 17 00:00:00 2001
From: Geeta Gharpure <geetagh@amazon.com>
Date: Tue, 13 Jun 2023 18:33:03 +0000
Subject: [PATCH] Enable test to verify membership recovery from backend

Signed-off-by: Geeta Gharpure <geetagh@amazon.com>
---
 .../api/membership/membership_test.go         | 65 +++++++++++++------
 1 file changed, 46 insertions(+), 19 deletions(-)

diff --git a/server/etcdserver/api/membership/membership_test.go b/server/etcdserver/api/membership/membership_test.go
index 728121e1c69..8e3a9691883 100644
--- a/server/etcdserver/api/membership/membership_test.go
+++ b/server/etcdserver/api/membership/membership_test.go
@@ -22,13 +22,14 @@ import (
 
 	"go.etcd.io/etcd/client/pkg/v3/types"
 	"go.etcd.io/etcd/server/v3/etcdserver/version"
+	serverversion "go.etcd.io/etcd/server/v3/etcdserver/version"
 
 	"go.uber.org/zap"
 )
 
 func TestAddRemoveMember(t *testing.T) {
 	c := newTestCluster(t, nil)
-	be := &backendMock{}
+	be := newMembershipBackend()
 	c.SetBackend(be)
 	c.AddMember(newTestMemberAsLearner(17, nil, "node17", nil), true)
 	c.RemoveMember(17, true)
@@ -39,34 +40,60 @@ func TestAddRemoveMember(t *testing.T) {
 	c.RemoveMember(17, true)
 	c.RemoveMember(18, true)
 
-	if false {
-		// TODO: Enable this code when Recover is reading membership from the backend.
-		c2 := newTestCluster(t, nil)
-		c2.SetBackend(be)
-		c2.Recover(func(*zap.Logger, *semver.Version) {})
-		assert.Equal(t, []*Member{{ID: types.ID(18),
-			Attributes: Attributes{Name: "node18"}}}, c2.Members())
-		assert.Equal(t, true, c2.IsIDRemoved(17))
-		assert.Equal(t, false, c2.IsIDRemoved(18))
-	}
+	c.AddMember(newTestMember(19, nil, "node19", nil), true)
+
+	// Recover from backend
+	c2 := newTestCluster(t, nil)
+	c2.SetBackend(be)
+	c2.Recover(func(*zap.Logger, *semver.Version) {})
+	assert.Equal(t, []*Member{{ID: types.ID(19),
+		Attributes: Attributes{Name: "node19"}}}, c2.Members())
+	assert.Equal(t, true, c2.IsIDRemoved(17))
+	assert.Equal(t, true, c2.IsIDRemoved(18))
+	assert.Equal(t, false, c2.IsIDRemoved(19))
 }
 
 type backendMock struct {
+	members       map[types.ID]*Member
+	removed       map[types.ID]bool
+	version       *semver.Version
+	downgradeInfo *version.DowngradeInfo
 }
 
 var _ MembershipBackend = (*backendMock)(nil)
 
+func newMembershipBackend() MembershipBackend {
+	return &backendMock{
+		members:       make(map[types.ID]*Member),
+		removed:       make(map[types.ID]bool),
+		downgradeInfo: &serverversion.DowngradeInfo{Enabled: false},
+	}
+}
+
 func (b *backendMock) MustCreateBackendBuckets() {}
 
-func (b *backendMock) ClusterVersionFromBackend() *semver.Version              { return nil }
-func (b *backendMock) MustSaveClusterVersionToBackend(version *semver.Version) {}
+func (b *backendMock) ClusterVersionFromBackend() *semver.Version { return b.version }
+func (b *backendMock) MustSaveClusterVersionToBackend(version *semver.Version) {
+	b.version = version
+}
 
 func (b *backendMock) MustReadMembersFromBackend() (x map[types.ID]*Member, y map[types.ID]bool) {
-	return
+	return b.members, b.removed
+}
+func (b *backendMock) MustSaveMemberToBackend(m *Member) {
+	b.members[m.ID] = m
+}
+func (b *backendMock) TrimMembershipFromBackend() error {
+	b.members = make(map[types.ID]*Member)
+	b.removed = make(map[types.ID]bool)
+	return nil
+}
+func (b *backendMock) MustDeleteMemberFromBackend(id types.ID) {
+	delete(b.members, id)
+	b.removed[id] = true
 }
-func (b *backendMock) MustSaveMemberToBackend(*Member)      {}
-func (b *backendMock) TrimMembershipFromBackend() error     { return nil }
-func (b *backendMock) MustDeleteMemberFromBackend(types.ID) {}
 
-func (b *backendMock) MustSaveDowngradeToBackend(*version.DowngradeInfo) {}
-func (b *backendMock) DowngradeInfoFromBackend() *version.DowngradeInfo  { return nil }
+func (b *backendMock) MustSaveDowngradeToBackend(downgradeInfo *version.DowngradeInfo) {
+	b.downgradeInfo = downgradeInfo
+}
+func (b *backendMock) DowngradeInfoFromBackend() *version.DowngradeInfo { return b.downgradeInfo }