From a8525f831580580df029003d79c35a37d2e6fc15 Mon Sep 17 00:00:00 2001
From: Dhia Ayachi <dhia@hashicorp.com>
Date: Wed, 15 Mar 2023 15:39:41 -0400
Subject: [PATCH] add snapshot restore test

---
 .../libs/cluster/container.go                 | 10 ++-
 .../test/snapshot/snapshot_restore_test.go    | 82 +++++++++++++++++++
 2 files changed, 90 insertions(+), 2 deletions(-)
 create mode 100644 test/integration/consul-container/test/snapshot/snapshot_restore_test.go

diff --git a/test/integration/consul-container/libs/cluster/container.go b/test/integration/consul-container/libs/cluster/container.go
index 3ad1b2d351..4e3e1c3990 100644
--- a/test/integration/consul-container/libs/cluster/container.go
+++ b/test/integration/consul-container/libs/cluster/container.go
@@ -154,11 +154,17 @@ func NewConsulContainer(ctx context.Context, config Config, cluster *Cluster, po
 		info AgentInfo
 	)
 	if httpPort > 0 {
-		uri, err := podContainer.PortEndpoint(ctx, "8500", "http")
+		for i := 0; i < 10; i++ {
+			uri, err := podContainer.PortEndpoint(ctx, "8500", "http")
+			if err != nil {
+				time.Sleep(500 * time.Millisecond)
+				continue
+			}
+			clientAddr = uri
+		}
 		if err != nil {
 			return nil, err
 		}
-		clientAddr = uri
 
 	} else if httpsPort > 0 {
 		uri, err := podContainer.PortEndpoint(ctx, "8501", "https")
diff --git a/test/integration/consul-container/test/snapshot/snapshot_restore_test.go b/test/integration/consul-container/test/snapshot/snapshot_restore_test.go
new file mode 100644
index 0000000000..7681ce97a5
--- /dev/null
+++ b/test/integration/consul-container/test/snapshot/snapshot_restore_test.go
@@ -0,0 +1,82 @@
+package snapshot
+
+import (
+	"fmt"
+	"github.com/hashicorp/consul/api"
+	"github.com/hashicorp/consul/sdk/testutil"
+	libcluster "github.com/hashicorp/consul/test/integration/consul-container/libs/cluster"
+	libtopology "github.com/hashicorp/consul/test/integration/consul-container/libs/topology"
+	"github.com/hashicorp/consul/test/integration/consul-container/libs/utils"
+	"github.com/stretchr/testify/require"
+	"io"
+	"testing"
+)
+
+func TestSnapshotRestore(t *testing.T) {
+	const (
+		numServers = 3
+	)
+
+	// Create initial cluster
+	cluster, _, _ := libtopology.NewCluster(t, &libtopology.ClusterConfig{
+		NumServers: numServers,
+		NumClients: 0,
+		BuildOpts: &libcluster.BuildOptions{
+			Datacenter:      "dc1",
+			ConsulImageName: utils.LatestImageName,
+			ConsulVersion:   utils.LatestVersion,
+		},
+		ApplyDefaultProxySettings: true,
+	})
+
+	client := cluster.APIClient(0)
+	libcluster.WaitForLeader(t, cluster, client)
+	libcluster.WaitForMembers(t, client, 3)
+
+	testutil.RunStep(t, "Create Data", func(t *testing.T) {
+		for i := 0; i < 100; i++ {
+			_, err := client.KV().Put(&api.KVPair{Key: fmt.Sprintf("key-%d", i), Value: []byte(fmt.Sprintf("value-%d", i))}, nil)
+			require.NoError(t, err)
+		}
+	})
+
+	var snapshot io.ReadCloser
+	testutil.RunStep(t, "Save snapshot", func(t *testing.T) {
+		var err error
+		snapshot, _, err = client.Snapshot().Save(nil)
+		require.NoError(t, err)
+	})
+
+	err := cluster.Terminate()
+	require.NoError(t, err)
+	// Create a fresh cluster from scratch
+	cluster2, _, _ := libtopology.NewCluster(t, &libtopology.ClusterConfig{
+		NumServers: numServers,
+		NumClients: 0,
+		BuildOpts: &libcluster.BuildOptions{
+			Datacenter:      "dc1",
+			ConsulImageName: utils.LatestImageName,
+			ConsulVersion:   utils.LatestVersion,
+		},
+		ApplyDefaultProxySettings: true,
+	})
+	client2 := cluster2.APIClient(0)
+
+	testutil.RunStep(t, "Restore saved snapshot", func(t *testing.T) {
+		libcluster.WaitForLeader(t, cluster2, client2)
+		libcluster.WaitForMembers(t, client2, 3)
+
+		// Restore the saved snapshot
+		require.NoError(t, client2.Snapshot().Restore(nil, snapshot))
+
+		libcluster.WaitForLeader(t, cluster2, client2)
+
+	})
+
+	for i := 0; i < 100; i++ {
+		kv, _, err := client2.KV().Get(fmt.Sprintf("key-%d", i), nil)
+		require.NoError(t, err)
+		require.Equal(t, kv.Key, fmt.Sprintf("key-%d", i))
+		require.Equal(t, kv.Value, []byte(fmt.Sprintf("value-%d", i)))
+	}
+}