From dec714a1c7024854a71fd00ff2fb3cee057ffb66 Mon Sep 17 00:00:00 2001
From: Mikkel Oscar Lyderik Larsen <mikkel.larsen@zalando.de>
Date: Thu, 30 Nov 2023 10:31:39 +0100
Subject: [PATCH] Improve the wait-for-apiserver ready check

Signed-off-by: Mikkel Oscar Lyderik Larsen <mikkel.larsen@zalando.de>
---
 provisioner/clusterpy.go      | 25 ++++++++++++++++------
 provisioner/clusterpy_test.go | 40 +++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/provisioner/clusterpy.go b/provisioner/clusterpy.go
index cf53e644..25a3f196 100644
--- a/provisioner/clusterpy.go
+++ b/provisioner/clusterpy.go
@@ -350,6 +350,7 @@ func (p *clusterpyProvisioner) Provision(ctx context.Context, logger *log.Entry,
 		cluster,
 		caNodePoolProvisioner,
 		karpenterProvisioner,
+		p.tokenSource,
 	)
 
 	err = nodePoolGroups["masters"].provisionNodePoolGroup(ctx, values, updater, cluster, p.applyOnly)
@@ -709,14 +710,26 @@ func (p *clusterpyProvisioner) removeEBSVolumes(awsAdapter *awsAdapter, cluster
 
 // waitForAPIServer waits a cluster API server to be ready. It's considered
 // ready when it's reachable.
-func waitForAPIServer(logger *log.Entry, server string, maxTimeout time.Duration) error {
+func waitForAPIServer(logger *log.Entry, cluster *api.Cluster, maxTimeout time.Duration, tokenSource oauth2.TokenSource) error {
 	logger.Infof("Waiting for API Server to be reachable")
 	client := &http.Client{}
 	timeout := time.Now().UTC().Add(maxTimeout)
 
 	for time.Now().UTC().Before(timeout) {
-		resp, err := client.Get(server)
-		if err == nil && resp.StatusCode < http.StatusInternalServerError {
+		req, err := http.NewRequest(http.MethodGet, cluster.APIServerURL+"/readyz", nil)
+		if err != nil {
+			return err
+		}
+
+		token, err := tokenSource.Token()
+		if err != nil {
+			return err
+		}
+
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
+
+		resp, err := client.Do(req)
+		if err == nil && resp.StatusCode == http.StatusOK {
 			return nil
 		}
 
@@ -725,7 +738,7 @@ func waitForAPIServer(logger *log.Entry, server string, maxTimeout time.Duration
 		time.Sleep(15 * time.Second)
 	}
 
-	return fmt.Errorf("'%s' was not ready after %s", server, maxTimeout.String())
+	return fmt.Errorf("%q was not ready after %s", cluster.APIServerURL, maxTimeout.String())
 }
 
 // setupAWSAdapter sets up the AWS Adapter used for communicating with AWS.
@@ -1136,7 +1149,7 @@ type nodePoolGroup struct {
 	ReadyFn     func() error
 }
 
-func groupNodePools(logger *log.Entry, cluster *api.Cluster, caProvisioner *AWSNodePoolProvisioner, karProvisioner *KarpenterNodePoolProvisioner) map[string]*nodePoolGroup {
+func groupNodePools(logger *log.Entry, cluster *api.Cluster, caProvisioner *AWSNodePoolProvisioner, karProvisioner *KarpenterNodePoolProvisioner, tokenSource oauth2.TokenSource) map[string]*nodePoolGroup {
 
 	var masters, workers, karpenterPools []*api.NodePool
 	for _, nodePool := range cluster.NodePools {
@@ -1157,7 +1170,7 @@ func groupNodePools(logger *log.Entry, cluster *api.Cluster, caProvisioner *AWSN
 			NodePools:   masters,
 			Provisioner: caProvisioner,
 			ReadyFn: func() error {
-				return waitForAPIServer(logger, cluster.APIServerURL, 15*time.Minute)
+				return waitForAPIServer(logger, cluster, 15*time.Minute, tokenSource)
 			},
 		},
 		"workers": {
diff --git a/provisioner/clusterpy_test.go b/provisioner/clusterpy_test.go
index 13671862..edc34d71 100644
--- a/provisioner/clusterpy_test.go
+++ b/provisioner/clusterpy_test.go
@@ -3,16 +3,21 @@ package provisioner
 import (
 	"errors"
 	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
+	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/zalando-incubator/cluster-lifecycle-manager/api"
 	"github.com/zalando-incubator/cluster-lifecycle-manager/channel"
 	"github.com/zalando-incubator/cluster-lifecycle-manager/pkg/kubernetes"
+	"golang.org/x/oauth2"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -477,3 +482,38 @@ second: 2
 		})
 	}
 }
+
+func TestWaitForAPIServer(t *testing.T) {
+	for _, tc := range []struct {
+		name         string
+		responseCode int
+	}{
+		{
+			name:         "test successful response from wait",
+			responseCode: 200,
+		},
+		{
+			name:         "test unuccessful response during wait",
+			responseCode: 500,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(tc.responseCode)
+			}))
+			defer ts.Close()
+
+			cluster := &api.Cluster{
+				APIServerURL: ts.URL,
+			}
+
+			tokenSource := oauth2.StaticTokenSource(&oauth2.Token{})
+			err := waitForAPIServer(log.WithField("cluster", "test"), cluster, 1*time.Millisecond, tokenSource)
+			if tc.responseCode == http.StatusOK {
+				require.NoError(t, err)
+			} else {
+				require.Error(t, err)
+			}
+		})
+	}
+}