Add WithUndo to experiments package

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 {

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"
 )
 

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 {

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