Wait for successful Prometheus metric

e2e/spot_migrator_test.go

@@ -9,6 +9,9 @@ import (
 	cloudproviderfake "github.com/hsbc/cost-manager/pkg/cloudprovider/fake"
 	"github.com/hsbc/cost-manager/pkg/kubernetes"
 	"github.com/hsbc/cost-manager/pkg/test"
+	"github.com/prometheus/client_golang/api"
+	prometheusv1 "github.com/prometheus/client_golang/api/prometheus/v1"
+	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
 	policyv1 "k8s.io/api/policy/v1"
@@ -154,11 +157,46 @@ func TestSpotMigrator(t *testing.T) {
 		require.False(t, node.Spec.Unschedulable)
 	}
 
-	// Delete Node; typically this would be done by the node-controller:
+	// Delete Node; typically this would be done by the node-controller but we simulate it here:
 	// https://github.com/hsbc/cost-manager/blob/bf176ada100e19a765d276aee1a0a2d6038275e0/pkg/controller/spot_migrator.go#L242-L250
 	err = kubeClient.Delete(ctx, node)
 	require.Nil(t, err)
 
+	// Wait for Prometheus metric to indicate successful migration
+	t.Logf("Waiting for successful Prometheus metric...")
+	pod, err := kubernetes.WaitForAnyReadyPod(ctx, kubeClient, client.InNamespace("monitoring"), client.MatchingLabels{"app.kubernetes.io/name": "prometheus"})
+	require.Nil(t, err)
+	// Port forward to Prometheus in the background
+	forwardedPort, close, err := kubernetes.PortForward(ctx, config.GetConfigOrDie(), pod.Namespace, pod.Name, 9090)
+	require.Nil(t, err)
+	defer func() {
+		err := close()
+		require.Nil(t, err)
+	}()
+	// Setup Prometheus client using local port forwarded port
+	prometheusAddress := fmt.Sprintf("http://127.0.0.1:%d", forwardedPort)
+	prometheusClient, err := api.NewClient(api.Config{
+		Address: prometheusAddress,
+	})
+	prometheusAPI := prometheusv1.NewAPI(prometheusClient)
+	for {
+		results, _, err := prometheusAPI.Query(ctx, "cost_manager_spot_migrator_operation_success_total", time.Now())
+		require.Nil(t, err)
+		// Any result with a value greater than 0 indicates migration success
+		migrationSuccess := false
+		for _, result := range results.(model.Vector) {
+			if result.Value > 0 {
+				migrationSuccess = true
+				break
+			}
+		}
+		if migrationSuccess {
+			break
+		}
+		time.Sleep(time.Second)
+	}
+	t.Logf("Found successful Prometheus metric!")
+
 	// Delete Namespace
 	err = kubeClient.Delete(ctx, namespace)
 	require.Nil(t, err)

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.17.0
+	github.com/prometheus/common v0.45.0
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/stretchr/testify v1.8.4
 	google.golang.org/api v0.149.0
@@ -86,7 +87,6 @@ require (
 	github.com/peterbourgon/diskv v2.0.1+incompatible // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/prometheus/client_model v0.4.1-0.20230718164431-9a2bf3000d16 // indirect
-	github.com/prometheus/common v0.45.0 // indirect
 	github.com/prometheus/procfs v0.11.1 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/spf13/cobra v1.7.0 // indirect

go.sum

@@ -169,6 +169,8 @@ github.com/jonboulle/clockwork v0.2.2 h1:UOGuzwb1PwsrDAObMuhUnj0p5ULPj8V/xJ7Kx9q
 github.com/jonboulle/clockwork v0.2.2/go.mod h1:Pkfl5aHPm1nk2H9h0bjmnJD/BcgbGXUBGnn1kMkgxc8=
 github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY=
 github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
+github.com/jpillora/backoff v1.0.0 h1:uvFg412JmmHBHw7iwprIxkPMI+sGQ4kzOWsMeHnm2EA=
+github.com/jpillora/backoff v1.0.0/go.mod h1:J/6gKK9jxlEcS3zixgDgUAsiuZ7yrSoa/FX5e0EB2j4=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
@@ -203,6 +205,8 @@ github.com/monochromegane/go-gitignore v0.0.0-20200626010858-205db1a8cc00 h1:n6/
 github.com/monochromegane/go-gitignore v0.0.0-20200626010858-205db1a8cc00/go.mod h1:Pm3mSP3c5uWn86xMLZ5Sa7JB9GsEZySvHYXCTK4E9q4=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
+github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f h1:KUppIJq7/+SVif2QVs3tOP0zanoHgBEVAwHxUSIzRqU=
+github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f h1:y5//uYreIhSUg3J1GEMiLbxo1LJaP8RfCpH6pymGZus=
 github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f/go.mod h1:ZdcZmHo+o7JKHSa8/e818NopupXU1YMK5fe1lsApnBw=
 github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78=

pkg/kubernetes/pod.go

@@ -0,0 +1,27 @@
+package kubernetes
+
+import (
+	"context"
+
+	corev1 "k8s.io/api/core/v1"
+	apiwatch "k8s.io/apimachinery/pkg/watch"
+	"k8s.io/client-go/tools/watch"
+	"k8s.io/kubectl/pkg/util/podutils"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+func WaitForAnyReadyPod(ctx context.Context, kubeClient client.WithWatch, opts ...client.ListOption) (*corev1.Pod, error) {
+	listerWatcher := NewListerWatcher(ctx, kubeClient, &corev1.PodList{}, opts...)
+	condition := func(event apiwatch.Event) (bool, error) {
+		pod, err := ParseWatchEventObject[*corev1.Pod](event)
+		if err != nil {
+			return false, err
+		}
+		return podutils.IsPodReady(pod), nil
+	}
+	event, err := watch.UntilWithSync(ctx, listerWatcher, &corev1.Pod{}, nil, condition)
+	if err != nil {
+		return nil, err
+	}
+	return event.Object.(*corev1.Pod), nil
+}