From 129caae0e4995b7861a6175195fab11c4f2c69b0 Mon Sep 17 00:00:00 2001
From: Otto Bittner <cobittner@posteo.net>
Date: Fri, 5 Aug 2022 17:23:17 +0200
Subject: [PATCH] AB#2258: Fix flacky retry_test.go

Co-authored-by: <mp@edgeless.systems>
Co-authored-by: <pm@edgeless.systems>
---
 internal/retry/retry_test.go     | 23 ++++++++++++++++-------
 internal/sigstore/verify_test.go |  5 +++++
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/internal/retry/retry_test.go b/internal/retry/retry_test.go
index dd5a88477..e31455ff4 100644
--- a/internal/retry/retry_test.go
+++ b/internal/retry/retry_test.go
@@ -7,9 +7,14 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"go.uber.org/goleak"
 	testclock "k8s.io/utils/clock/testing"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+
 func TestDo(t *testing.T) {
 	testCases := map[string]struct {
 		cancel  bool
@@ -64,16 +69,15 @@ func TestDo(t *testing.T) {
 			defer cancel()
 
 			go func() { retrierResult <- retrier.Do(ctx) }()
-
-			if tc.cancel {
-				cancel()
-			}
-
 			for _, err := range tc.errors {
 				doer.errC <- err
 				clock.Step(retrier.interval)
 			}
 
+			if tc.cancel {
+				cancel()
+			}
+
 			assert.Equal(tc.wantErr, <-retrierResult)
 		})
 	}
@@ -89,8 +93,13 @@ func newStubDoer() *stubDoer {
 	}
 }
 
-func (d *stubDoer) Do(_ context.Context) error {
-	return <-d.errC
+func (d *stubDoer) Do(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-d.errC:
+		return err
+	}
 }
 
 func isRetriable(err error) bool {
diff --git a/internal/sigstore/verify_test.go b/internal/sigstore/verify_test.go
index b9b279bd4..0eff1f957 100644
--- a/internal/sigstore/verify_test.go
+++ b/internal/sigstore/verify_test.go
@@ -4,8 +4,13 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+
 func TestVerifySignature(t *testing.T) {
 	testCases := map[string]struct {
 		content   []byte