From dcefd8acd3211e3efab4190491e02434fc967bd2 Mon Sep 17 00:00:00 2001
From: Shiming Zhang <wzshiming@hotmail.com>
Date: Tue, 3 Sep 2024 09:17:27 +0800
Subject: [PATCH] Fix flake test

---
 test/e2e/benchmark.go |  7 +++----
 test/e2e/workable.go  | 25 ++++++++++++++++++++-----
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/test/e2e/benchmark.go b/test/e2e/benchmark.go
index 3a2df89fb..968dab28f 100644
--- a/test/e2e/benchmark.go
+++ b/test/e2e/benchmark.go
@@ -45,13 +45,12 @@ func waitResource(ctx context.Context, t *testing.T, kwokctlPath, name, resource
 		if got == want {
 			return nil
 		}
-		if got == prev {
-			t.Logf("Error %s %d not changed\n", resource, got)
+		all := strings.Count(raw, "\n")
+		t.Logf("%s %d/%d => %d\n", resource, got, all, want)
+		if prev != 0 && got == prev {
 			return fmt.Errorf("resource %s not changed", resource)
 		}
 		prev = got
-		all := strings.Count(raw, "\n")
-		t.Logf("%s %d/%d => %d\n", resource, got, all, want)
 		if gap != 0 && got != 0 && (all-got) > gap {
 			if tolerance > 0 {
 				t.Logf("Error %s gap too large, actual: %d, expected: %d, retrying...\n", resource, all-got, gap)
diff --git a/test/e2e/workable.go b/test/e2e/workable.go
index 00857e32f..fede2d6d1 100644
--- a/test/e2e/workable.go
+++ b/test/e2e/workable.go
@@ -27,7 +27,9 @@ import (
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
+	"sigs.k8s.io/e2e-framework/klient/wait"
 	"sigs.k8s.io/e2e-framework/pkg/envconf"
 	"sigs.k8s.io/e2e-framework/pkg/features"
 )
@@ -49,14 +51,27 @@ func testWorkable(ctx context.Context, kwokctlPath, name string) error {
 	if err != nil {
 		return fmt.Errorf("failed to scale pod: %w", err)
 	}
-	output, err := exec.CommandContext(ctx, kwokctlPath, "--name", name, "kubectl", "get", "pod").Output()
+
+	err = wait.For(
+		func(ctx context.Context) (done bool, err error) {
+			output, err := exec.CommandContext(ctx, kwokctlPath, "--name", name, "kubectl", "get", "pod").Output()
+			if err != nil {
+				return false, err
+			}
+			// TODO: Use json output instead and check that node and pod work as expected
+			if !strings.Contains(string(output), "Running") {
+				return false, fmt.Errorf("pod not running")
+			}
+
+			return true, nil
+		},
+		wait.WithContext(ctx),
+		wait.WithTimeout(10*time.Second),
+	)
 	if err != nil {
 		return err
 	}
-	// TODO: Use json output instead and check that node and pod work as expected
-	if !strings.Contains(string(output), "Running") {
-		return fmt.Errorf("pod not running")
-	}
+
 	return nil
 }