From 6547f203e49783e2ca27e4a47d3299b6e402d6cb Mon Sep 17 00:00:00 2001
From: Aaron Liang <aaronliang@google.com>
Date: Wed, 5 Apr 2023 22:10:07 +0000
Subject: [PATCH] Add e2e test for leader election

---
 test/e2e/controller/crash_test.go | 62 ++++++++++++++++++++++++++++---
 1 file changed, 57 insertions(+), 5 deletions(-)

diff --git a/test/e2e/controller/crash_test.go b/test/e2e/controller/crash_test.go
index e30f78a483..297310944b 100644
--- a/test/e2e/controller/crash_test.go
+++ b/test/e2e/controller/crash_test.go
@@ -22,6 +22,7 @@ import (
 	agonesv1 "agones.dev/agones/pkg/apis/agones/v1"
 	e2eframework "agones.dev/agones/test/e2e/framework"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
@@ -55,7 +56,48 @@ func TestGameServerUnhealthyAfterDeletingPodWhileControllerDown(t *testing.T) {
 	_, err = framework.WaitForGameServerState(t, readyGs, agonesv1.GameServerStateUnhealthy, 3*time.Minute)
 	assert.NoError(t, err)
 	logger.Info("waiting for Agones controller to come back to running")
-	assert.NoError(t, waitForAgonesControllerRunning(ctx))
+	assert.NoError(t, waitForAgonesControllerRunning(ctx, -1))
+}
+
+func TestLeaderElectionAfterDeletingLeader(t *testing.T) {
+	logger := e2eframework.TestLogger(t)
+	gs := framework.DefaultGameServer(defaultNs)
+	ctx := context.Background()
+
+	// if !runtime.FeatureEnabled(runtime.FeatureSplitControllerAndExtensions) {
+	// 	t.Skip("Skip test. SplitControllerAndExtensions feature is not enabled")
+	// }
+
+	err := waitForAgonesControllerRunning(ctx, -1)
+	require.NoError(t, err, "Could not ensure controller running")
+
+	list, err := getAgonesControllerPods(ctx)
+	require.NoError(t, err, "Could not get list of Extension pods")
+	if len(list.Items) == 1 {
+		t.Skip("Skip test. Leader Election is not enabled since there is only 1 controller")
+	}
+
+	replication := len(list.Items)
+
+	// Deleting one of the pods which cause the other one to most definitely be leader
+	willBeLeader := list.Items[0].ObjectMeta.Name
+	for _, pod := range list.Items[1:] {
+		err = deleteAgonesControllerPod(ctx, pod.ObjectMeta.Name)
+		require.NoError(t, err, "Could not delete controller pod")
+	}
+
+	err = waitForAgonesControllerRunning(ctx, replication)
+	require.NoError(t, err, "Could not get controller ready after delete")
+
+	err = deleteAgonesControllerPod(ctx, willBeLeader)
+	require.NoError(t, err, "Could not delete leader controller pod")
+
+	// Do something that will test controller like make game server
+	readyGs, err := framework.CreateGameServerAndWaitUntilReady(t, defaultNs, gs)
+	if err != nil {
+		t.Fatalf("Could not get a GameServer ready: %v", err)
+	}
+	logger.WithField("gsKey", readyGs.ObjectMeta.Name).Info("GameServer Ready")
 }
 
 // deleteAgonesControllerPods deletes all the Controller pods for the Agones controller,
@@ -66,10 +108,8 @@ func deleteAgonesControllerPods(ctx context.Context) error {
 		return err
 	}
 
-	policy := metav1.DeletePropagationBackground
 	for i := range list.Items {
-		err = framework.KubeClient.CoreV1().Pods("agones-system").Delete(ctx, list.Items[i].ObjectMeta.Name,
-			metav1.DeleteOptions{PropagationPolicy: &policy})
+		err = deleteAgonesControllerPod(ctx, list.Items[i].ObjectMeta.Name)
 		if err != nil {
 			return err
 		}
@@ -77,13 +117,17 @@ func deleteAgonesControllerPods(ctx context.Context) error {
 	return nil
 }
 
-func waitForAgonesControllerRunning(ctx context.Context) error {
+func waitForAgonesControllerRunning(ctx context.Context, wantReplicas int) error {
 	return wait.PollImmediate(time.Second, 5*time.Minute, func() (bool, error) {
 		list, err := getAgonesControllerPods(ctx)
 		if err != nil {
 			return true, err
 		}
 
+		if wantReplicas != -1 && len(list.Items) != wantReplicas {
+			return false, nil
+		}
+
 		for i := range list.Items {
 			for _, c := range list.Items[i].Status.ContainerStatuses {
 				if c.State.Running == nil {
@@ -101,3 +145,11 @@ func getAgonesControllerPods(ctx context.Context) (*corev1.PodList, error) {
 	opts := metav1.ListOptions{LabelSelector: labels.Set{"agones.dev/role": "controller"}.String()}
 	return framework.KubeClient.CoreV1().Pods("agones-system").List(ctx, opts)
 }
+
+// deleteAgonesControllerPod deletes a Agones controller pod
+func deleteAgonesControllerPod(ctx context.Context, podName string) error {
+	policy := metav1.DeletePropagationBackground
+	err := framework.KubeClient.CoreV1().Pods("agones-system").Delete(ctx, podName,
+		metav1.DeleteOptions{PropagationPolicy: &policy})
+	return err
+}