From 3275836e9b360bc84efb20c7e9804695be1d1a12 Mon Sep 17 00:00:00 2001
From: Ben Moskovitz <ben.m@buildkite.com>
Date: Wed, 3 May 2023 14:58:11 +1000
Subject: [PATCH] Add WithUndo to experiments package

---
 .../integration/checkout_integration_test.go   | 18 +++---------------
 .../integration/hooks_integration_test.go      |  1 +
 .../integration/job_api_integration_test.go    |  2 +-
 experiments/experiments.go                     |  5 +++++
 4 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/bootstrap/integration/checkout_integration_test.go b/bootstrap/integration/checkout_integration_test.go
index 2f3eb107ae..08cf6a149e 100644
--- a/bootstrap/integration/checkout_integration_test.go
+++ b/bootstrap/integration/checkout_integration_test.go
@@ -26,21 +26,9 @@ var commitPattern = bintest.MatchPattern(`(?ms)\Acommit [0-9a-f]+\nabbrev-commit
 // We expect this arg multiple times, just define it once.
 var gitShowFormatArg = "--format=commit %H%nabbrev-commit %h%nAuthor: %an <%ae>%n%n%w(0,4,4)%B"
 
-// Enable an experiment, returning a function to restore the previous state.
-// Usage: defer experimentWithUndo("foo")()
-func experimentWithUndo(name string) func() {
-	prev := experiments.IsEnabled(name)
-	experiments.Enable(name)
-	return func() {
-		if !prev {
-			experiments.Disable(name)
-		}
-	}
-}
-
 func TestCheckingOutGitHubPullRequestsWithGitMirrorsExperiment(t *testing.T) {
 	// t.Parallel() cannot be used with experiments.Enable()
-	defer experimentWithUndo(experiments.GitMirrors)()
+	defer experiments.WithUndo(experiments.GitMirrors)()
 
 	tester, err := NewBootstrapTester()
 	if err != nil {
@@ -81,7 +69,7 @@ func TestCheckingOutGitHubPullRequestsWithGitMirrorsExperiment(t *testing.T) {
 
 func TestWithResolvingCommitExperiment(t *testing.T) {
 	// t.Parallel() cannot be used with experiments.Enable()
-	defer experimentWithUndo(experiments.ResolveCommitAfterCheckout)()
+	defer experiments.WithUndo(experiments.ResolveCommitAfterCheckout)()
 
 	tester, err := NewBootstrapTester()
 	if err != nil {
@@ -691,7 +679,7 @@ func TestRepositorylessCheckout(t *testing.T) {
 
 func TestGitMirrorEnv(t *testing.T) {
 	// t.Parallel() cannot test experiment flags in parallel
-	defer experimentWithUndo(experiments.GitMirrors)()
+	defer experiments.WithUndo(experiments.GitMirrors)()
 
 	tester, err := NewBootstrapTester()
 	if err != nil {
diff --git a/bootstrap/integration/hooks_integration_test.go b/bootstrap/integration/hooks_integration_test.go
index 119aa872d4..4bf6aefd6a 100644
--- a/bootstrap/integration/hooks_integration_test.go
+++ b/bootstrap/integration/hooks_integration_test.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/buildkite/agent/v3/bootstrap/shell"
+	"github.com/buildkite/agent/v3/experiments"
 	"github.com/buildkite/bintest/v3"
 )
 
diff --git a/bootstrap/integration/job_api_integration_test.go b/bootstrap/integration/job_api_integration_test.go
index 554e06447d..1dd8a34fb6 100644
--- a/bootstrap/integration/job_api_integration_test.go
+++ b/bootstrap/integration/job_api_integration_test.go
@@ -16,7 +16,7 @@ import (
 )
 
 func TestBootstrapRunsJobAPI(t *testing.T) {
-	defer experimentWithUndo(experiments.JobAPI)()
+	defer experiments.WithUndo(experiments.JobAPI)()
 
 	tester, err := NewBootstrapTester()
 	if err != nil {
diff --git a/experiments/experiments.go b/experiments/experiments.go
index a2e90e0c80..40de7d7ca8 100644
--- a/experiments/experiments.go
+++ b/experiments/experiments.go
@@ -32,6 +32,11 @@ var (
 	experiments = make(map[string]bool, len(Available))
 )
 
+func WithUndo(key string) func() {
+	Enable(key)
+	return func() { Disable(key) }
+}
+
 // Enable a particular experiment in the agent.
 func Enable(key string) (known bool) {
 	experiments[key] = true