chore: add testutil.Eventually and friends (#3389)

johnstcn · mafredri · web-flow · commit 01fe5e668e80 · 2022-08-05T16:34:44.000+01:00
This PR adds a `testutil` function aimed to replace `require.Eventually`.

Before:
```go
require.Eventually(t, func() bool { ... }, testutil.WaitShort, testutil.IntervalFast)
```

After:
```go
require.True(t, testutil.EventuallyShort(t, func(ctx context.Context) bool { ... }))

// or the full incantation if you need more control
ctx, cancel := context.WithTimeout(ctx.Background(), testutil.WaitLong)
require.True(t, testutil.Eventually(t, ctx, func(ctx context.Context) bool { ... }, testutil.IntervalSlow))
```

Co-authored-by: Mathias Fredriksson &lt;mafredri@gmail.com&gt;
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -411,11 +411,11 @@ func AwaitTemplateVersionJob(t *testing.T, client *codersdk.Client, version uuid
 
 	t.Logf("waiting for template version job %s", version)
 	var templateVersion codersdk.TemplateVersion
-	require.Eventually(t, func() bool {
+	require.True(t, testutil.EventuallyShort(t, func(ctx context.Context) bool {
 		var err error
-		templateVersion, err = client.TemplateVersion(context.Background(), version)
+		templateVersion, err = client.TemplateVersion(ctx, version)
 		return assert.NoError(t, err) && templateVersion.Job.CompletedAt != nil
-	}, testutil.WaitShort, testutil.IntervalFast)
+	}))
 	return templateVersion
 }
 
@@ -425,11 +425,10 @@ func AwaitWorkspaceBuildJob(t *testing.T, client *codersdk.Client, build uuid.UU
 
 	t.Logf("waiting for workspace build job %s", build)
 	var workspaceBuild codersdk.WorkspaceBuild
-	require.Eventually(t, func() bool {
-		var err error
-		workspaceBuild, err = client.WorkspaceBuild(context.Background(), build)
+	require.True(t, testutil.EventuallyShort(t, func(ctx context.Context) bool {
+		workspaceBuild, err := client.WorkspaceBuild(ctx, build)
 		return assert.NoError(t, err) && workspaceBuild.Job.CompletedAt != nil
-	}, testutil.WaitShort, testutil.IntervalFast)
+	}))
 	return workspaceBuild
 }
 
@@ -439,21 +438,22 @@ func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, build uuid.UUID
 
 	t.Logf("waiting for workspace agents (build %s)", build)
 	var resources []codersdk.WorkspaceResource
-	require.Eventually(t, func() bool {
+	require.True(t, testutil.EventuallyLong(t, func(ctx context.Context) bool {
 		var err error
-		resources, err = client.WorkspaceResourcesByBuild(context.Background(), build)
+		resources, err = client.WorkspaceResourcesByBuild(ctx, build)
 		if !assert.NoError(t, err) {
 			return false
 		}
 		for _, resource := range resources {
 			for _, agent := range resource.Agents {
 				if agent.Status != codersdk.WorkspaceAgentConnected {
+					t.Logf("agent %s not connected yet", agent.Name)
 					return false
 				}
 			}
 		}
 		return true
-	}, testutil.WaitLong, testutil.IntervalMedium)
+	}))
 	return resources
 }
 
diff --git a/coderd/coderdtest/coderdtest_test.go b/coderd/coderdtest/coderdtest_test.go
@@ -19,7 +19,7 @@ func TestNew(t *testing.T) {
 	})
 	user := coderdtest.CreateFirstUser(t, client)
 	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
-	coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+	_ = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
 	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
 	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
 	coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
diff --git a/scripts/rules.go b/scripts/rules.go
@@ -91,7 +91,7 @@ func useStandardTimeoutsAndDelaysInTests(m dsl.Matcher) {
 	m.Import("github.com/coder/coder/testutil")
 
 	m.Match(`context.WithTimeout($ctx, $duration)`).
-		Where(m.File().Imports("testing") && !m["duration"].Text.Matches("^testutil\\.")).
+		Where(m.File().Imports("testing") && !m.File().PkgPath.Matches("testutil$") && !m["duration"].Text.Matches("^testutil\\.")).
 		At(m["duration"]).
 		Report("Do not use magic numbers in test timeouts and delays. Use the standard testutil.Wait* or testutil.Interval* constants instead.")
 
diff --git a/testutil/eventually.go b/testutil/eventually.go
@@ -0,0 +1,68 @@
+package testutil
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// Eventually is like require.Eventually except it allows passing
+// a context into the condition. It is safe to use with `require.*`.
+//
+// If ctx times out, the test will fail, but not immediately.
+// It is the caller's responsibility to exit early if required.
+//
+// It is the caller's responsibility to ensure that ctx has a
+// deadline or timeout set. Eventually will panic if this is not
+// the case in order to avoid potentially waiting forever.
+//
+// condition is not run in a goroutine; use the provided
+// context argument for cancellation if required.
+func Eventually(ctx context.Context, t testing.TB, condition func(context.Context) bool, tick time.Duration) bool {
+	t.Helper()
+
+	if _, ok := ctx.Deadline(); !ok {
+		panic("developer error: must set deadline or timeout on ctx")
+	}
+
+	ticker := time.NewTicker(tick)
+	defer ticker.Stop()
+	for tick := ticker.C; ; {
+		select {
+		case <-ctx.Done():
+			assert.NoError(t, ctx.Err(), "Eventually timed out")
+			return false
+		case <-tick:
+			assert.NoError(t, ctx.Err(), "Eventually timed out")
+			if condition(ctx) {
+				return true
+			}
+		}
+	}
+}
+
+// EventuallyShort is a convenience function that runs Eventually with
+// IntervalFast and times out after WaitShort.
+func EventuallyShort(t testing.TB, condition func(context.Context) bool) bool {
+	ctx, cancel := context.WithTimeout(context.Background(), WaitShort)
+	defer cancel()
+	return Eventually(ctx, t, condition, IntervalFast)
+}
+
+// EventuallyMedium is a convenience function that runs Eventually with
+// IntervalMedium and times out after WaitMedium.
+func EventuallyMedium(t testing.TB, condition func(context.Context) bool) bool {
+	ctx, cancel := context.WithTimeout(context.Background(), WaitMedium)
+	defer cancel()
+	return Eventually(ctx, t, condition, IntervalMedium)
+}
+
+// EventuallyLong is a convenience function that runs Eventually with
+// IntervalSlow and times out after WaitLong.
+func EventuallyLong(t testing.TB, condition func(context.Context) bool) bool {
+	ctx, cancel := context.WithTimeout(context.Background(), WaitLong)
+	defer cancel()
+	return Eventually(ctx, t, condition, IntervalSlow)
+}
diff --git a/testutil/eventually_test.go b/testutil/eventually_test.go
@@ -0,0 +1,70 @@
+package testutil_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/goleak"
+
+	"github.com/coder/coder/testutil"
+)
+
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+
+func TestEventually(t *testing.T) {
+	t.Parallel()
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+		state := 0
+		condition := func(_ context.Context) bool {
+			defer func() {
+				state++
+			}()
+			return state > 2
+		}
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		defer cancel()
+		testutil.Eventually(ctx, t, condition, testutil.IntervalFast)
+	})
+
+	t.Run("Timeout", func(t *testing.T) {
+		t.Parallel()
+		condition := func(_ context.Context) bool {
+			return false
+		}
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		defer cancel()
+		mockT := new(testing.T)
+		testutil.Eventually(ctx, mockT, condition, testutil.IntervalFast)
+		assert.True(t, mockT.Failed())
+	})
+
+	t.Run("Panic", func(t *testing.T) {
+		t.Parallel()
+
+		panicky := func() {
+			mockT := new(testing.T)
+			condition := func(_ context.Context) bool { return true }
+			testutil.Eventually(context.Background(), mockT, condition, testutil.IntervalFast)
+		}
+		assert.Panics(t, panicky)
+	})
+
+	t.Run("Short", func(t *testing.T) {
+		t.Parallel()
+		testutil.EventuallyShort(t, func(_ context.Context) bool { return true })
+	})
+
+	t.Run("Medium", func(t *testing.T) {
+		t.Parallel()
+		testutil.EventuallyMedium(t, func(_ context.Context) bool { return true })
+	})
+
+	t.Run("Long", func(t *testing.T) {
+		t.Parallel()
+		testutil.EventuallyLong(t, func(_ context.Context) bool { return true })
+	})
+}
