Replace os.Setenv with testing.SetEnv

8 files changed, 14 insertions, 35 deletions

diff --git a/internal/feature/feature_test.go b/internal/feature/feature_test.go
index 88de8419..fec4e2a2 100644
--- a/internal/feature/feature_test.go
+++ b/internal/feature/feature_test.go
@@ -4,8 +4,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
-
-	"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
 )
 
 func TestEnabled(t *testing.T) {
@@ -41,7 +39,7 @@ func TestEnabled(t *testing.T) {
 				EnvVariable:    "testFeatureFlag",
 				defaultEnabled: tt.defaultEnabled,
 			}
-			testhelpers.SetEnvironmentVariable(t, feature.EnvVariable, tt.envVal)
+			t.Setenv(feature.EnvVariable, tt.envVal)
 			require.Equal(t, tt.expected, feature.Enabled())
 		})
 	}
diff --git a/internal/handlers/ratelimiter_test.go b/internal/handlers/ratelimiter_test.go
index 43acfc9a..1ff1c0e5 100644
--- a/internal/handlers/ratelimiter_test.go
+++ b/internal/handlers/ratelimiter_test.go
@@ -3,6 +3,7 @@ package handlers
 import (
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -75,8 +76,8 @@ func TestRatelimiter(t *testing.T) {
 
 	for name, tc := range tt {
 		t.Run(name, func(t *testing.T) {
-			testhelpers.StubFeatureFlagValue(t, feature.EnforceIPRateLimits.EnvVariable, tc.sourceIPEnforced)
-			testhelpers.StubFeatureFlagValue(t, feature.EnforceDomainRateLimits.EnvVariable, tc.domainEnforced)
+			t.Setenv(feature.EnforceIPRateLimits.EnvVariable, strconv.FormatBool(tc.sourceIPEnforced))
+			t.Setenv(feature.EnforceDomainRateLimits.EnvVariable, strconv.FormatBool(tc.domainEnforced))
 
 			conf := config.RateLimit{
 				SourceIPLimitPerSecond: 0.1,
diff --git a/internal/ratelimiter/middleware_test.go b/internal/ratelimiter/middleware_test.go
index 25ac08b5..167a9e68 100644
--- a/internal/ratelimiter/middleware_test.go
+++ b/internal/ratelimiter/middleware_test.go
@@ -77,7 +77,7 @@ func TestMiddlewareDenyRequestsAfterBurst(t *testing.T) {
 
 	for tn, tc := range tcs {
 		t.Run(tn, func(t *testing.T) {
-			testhelpers.StubFeatureFlagValue(t, feature.EnforceIPRateLimits.EnvVariable, tc.enforce)
+			t.Setenv(feature.EnforceIPRateLimits.EnvVariable, strconv.FormatBool(tc.enforce))
 
 			rl := New(
 				"rate_limiter",
diff --git a/internal/redirects/matching_test.go b/internal/redirects/matching_test.go
index 6ebb12fb..1460dfb1 100644
--- a/internal/redirects/matching_test.go
+++ b/internal/redirects/matching_test.go
@@ -5,6 +5,8 @@ import (
 
 	"github.com/stretchr/testify/require"
 	netlifyRedirects "github.com/tj/go-redirects"
+
+	"gitlab.com/gitlab-org/gitlab-pages/internal/feature"
 )
 
 type testCaseData struct {
@@ -55,7 +57,7 @@ var testsWithoutPlaceholders = map[string]testCaseData{
 }
 
 func Test_matchesRule(t *testing.T) {
-	enablePlaceholders(t)
+	t.Setenv(feature.RedirectsPlaceholders.EnvVariable, "true")
 
 	tests := mergeTestSuites(testsWithoutPlaceholders, map[string]testCaseData{
 		// Note: the following 3 cases behave differently when
diff --git a/internal/redirects/redirects_benchmark_test.go b/internal/redirects/redirects_benchmark_test.go
index f02eb679..6f5d3dd2 100644
--- a/internal/redirects/redirects_benchmark_test.go
+++ b/internal/redirects/redirects_benchmark_test.go
@@ -55,7 +55,7 @@ func BenchmarkRedirectsRewrite_withoutPlaceholders(b *testing.B) {
 }
 
 func BenchmarkRedirectsRewrite_PlaceholdersEnabled(b *testing.B) {
-	enablePlaceholders(b)
+	b.Setenv(feature.RedirectsPlaceholders.EnvVariable, "true")
 
 	b.Run("10 redirects", func(b *testing.B) { benchmarkRedirectsRewrite(b, 10) })
 	b.Run("100 redirects", func(b *testing.B) { benchmarkRedirectsRewrite(b, 100) })
diff --git a/internal/redirects/redirects_test.go b/internal/redirects/redirects_test.go
index c0af60a7..66c5d733 100644
--- a/internal/redirects/redirects_test.go
+++ b/internal/redirects/redirects_test.go
@@ -16,13 +16,8 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
 )
 
-// enablePlaceholders enables redirect placeholders in tests
-func enablePlaceholders(t testing.TB) {
-	testhelpers.StubFeatureFlagValue(t, feature.RedirectsPlaceholders.EnvVariable, true)
-}
-
 func TestRedirectsRewrite(t *testing.T) {
-	enablePlaceholders(t)
+	t.Setenv(feature.RedirectsPlaceholders.EnvVariable, "true")
 
 	tests := []struct {
 		name           string
diff --git a/internal/redirects/validations_test.go b/internal/redirects/validations_test.go
index bd108e6c..6d6fbb3d 100644
--- a/internal/redirects/validations_test.go
+++ b/internal/redirects/validations_test.go
@@ -6,10 +6,12 @@ import (
 
 	"github.com/stretchr/testify/require"
 	netlifyRedirects "github.com/tj/go-redirects"
+
+	"gitlab.com/gitlab-org/gitlab-pages/internal/feature"
 )
 
 func TestRedirectsValidateUrl(t *testing.T) {
-	enablePlaceholders(t)
+	t.Setenv(feature.RedirectsPlaceholders.EnvVariable, "true")
 
 	tests := map[string]struct {
 		url         string
@@ -84,7 +86,7 @@ func TestRedirectsValidateUrlNoPlaceholders(t *testing.T) {
 }
 
 func TestRedirectsValidateRule(t *testing.T) {
-	enablePlaceholders(t)
+	t.Setenv(feature.RedirectsPlaceholders.EnvVariable, "true")
 
 	tests := map[string]struct {
 		rule        string
diff --git a/internal/testhelpers/testhelpers.go b/internal/testhelpers/testhelpers.go
index 2f708743..fdd3e7d4 100644
--- a/internal/testhelpers/testhelpers.go
+++ b/internal/testhelpers/testhelpers.go
@@ -7,7 +7,6 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"os"
-	"strconv"
 	"testing"
 
 	"github.com/sirupsen/logrus"
@@ -63,24 +62,6 @@ func Getwd(t *testing.T) string {
 	return wd
 }
 
-// SetEnvironmentVariable for testing, restoring the original value on t.Cleanup
-func SetEnvironmentVariable(t testing.TB, key, value string) {
-	t.Helper()
-
-	orig := os.Getenv(key)
-
-	err := os.Setenv(key, value)
-	require.NoError(t, err)
-
-	t.Cleanup(func() {
-		require.NoError(t, os.Setenv(key, orig))
-	})
-}
-
-func StubFeatureFlagValue(t testing.TB, envVar string, value bool) {
-	SetEnvironmentVariable(t, envVar, strconv.FormatBool(value))
-}
-
 func PerformRequest(t *testing.T, handler http.Handler, r *http.Request) (int, string) {
 	t.Helper()