From 4ec4587349ca977fcca3e537608814461f0a3b85 Mon Sep 17 00:00:00 2001
From: joey <zchengjoey@gmail.com>
Date: Wed, 18 Oct 2023 11:05:08 +0800
Subject: [PATCH] fix waiter test is flaky

Signed-off-by: chengjoey <zchengjoey@gmail.com>
---
 cmd/entrypoint/waiter.go            | 22 +++++++++++-----------
 pkg/entrypoint/entrypointer_test.go | 10 ++++------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/cmd/entrypoint/waiter.go b/cmd/entrypoint/waiter.go
index 729c747e33c..656b015af74 100644
--- a/cmd/entrypoint/waiter.go
+++ b/cmd/entrypoint/waiter.go
@@ -54,17 +54,6 @@ func (rw *realWaiter) Wait(ctx context.Context, file string, expectContent bool,
 		return nil
 	}
 	for {
-		select {
-		case <-ctx.Done():
-			if errors.Is(ctx.Err(), context.Canceled) {
-				return entrypoint.ErrContextCanceled
-			}
-			if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-				return entrypoint.ErrContextDeadlineExceeded
-			}
-			return nil
-		case <-time.After(rw.waitPollingInterval):
-		}
 		if info, err := os.Stat(file); err == nil {
 			if !expectContent || info.Size() > 0 {
 				return nil
@@ -84,6 +73,17 @@ func (rw *realWaiter) Wait(ctx context.Context, file string, expectContent bool,
 			}
 			return skipError("error file present, bail and skip the step")
 		}
+		select {
+		case <-ctx.Done():
+			if errors.Is(ctx.Err(), context.Canceled) {
+				return entrypoint.ErrContextCanceled
+			}
+			if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+				return entrypoint.ErrContextDeadlineExceeded
+			}
+			return nil
+		case <-time.After(rw.waitPollingInterval):
+		}
 	}
 }
 
diff --git a/pkg/entrypoint/entrypointer_test.go b/pkg/entrypoint/entrypointer_test.go
index d2b749cab1b..06113dbb628 100644
--- a/pkg/entrypoint/entrypointer_test.go
+++ b/pkg/entrypoint/entrypointer_test.go
@@ -196,19 +196,15 @@ func TestEntrypointer(t *testing.T) {
 			}
 
 			if len(c.waitFiles) > 0 {
-				fw.Lock()
 				if fw.waited == nil {
 					t.Error("Wanted waited file, got nil")
-				} else if !reflect.DeepEqual(fw.waited, append(c.waitFiles, "/tekton/downward/cancel")) {
+				} else if !reflect.DeepEqual(fw.waited, c.waitFiles) {
 					t.Errorf("Waited for %v, want %v", fw.waited, c.waitFiles)
 				}
-				fw.Unlock()
 			}
-			fw.Lock()
-			if len(c.waitFiles) == 0 && len(fw.waited) != 1 {
+			if len(c.waitFiles) == 0 && fw.waited != nil {
 				t.Errorf("Waited for file when not required")
 			}
-			fw.Unlock()
 
 			wantArgs := append([]string{c.entrypoint}, c.args...)
 			if len(wantArgs) != 0 {
@@ -699,6 +695,8 @@ type fakeWaiter struct {
 func (f *fakeWaiter) Wait(ctx context.Context, file string, _ bool, _ bool) error {
 	if file == pod.DownwardMountCancelFile && f.waitCancelDuration > 0 {
 		time.Sleep(f.waitCancelDuration)
+	} else if file == pod.DownwardMountCancelFile {
+		return nil
 	}
 	f.Lock()
 	f.waited = append(f.waited, file)