test: use T.Setenv to set env vars in tests (#1179)

This commit replaces `os.Setenv` with `t.Setenv` in tests. The
environment variable is automatically restored to its original value
when the test and all its subtests complete.

Reference: https://pkg.go.dev/testing#T.Setenv
Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Author: Juneezee

cmd/readiness/readiness_test.go

@@ -141,31 +141,27 @@ func TestNoDeadlockForImmediateWaitRs(t *testing.T) {
 // (as Agent doesn't marks all the step statuses finished when it reaches the goal) but this doesn't affect the result
 // as the whole plan is complete already
 func TestHeadlessAgentHasntReachedGoal(t *testing.T) {
-	_ = os.Setenv(headlessAgent, "true")
+	t.Setenv(headlessAgent, "true")
 	c := testConfig("testdata/health-status-ok.json")
 	c.ClientSet = fake.NewSimpleClientset(testdata.TestPod(c.Namespace, c.Hostname), testdata.TestSecret(c.Namespace, c.AutomationConfigSecretName, 6))
 	ready, err := isPodReady(c)
 	assert.False(t, ready)
 	assert.NoError(t, err)
 	thePod, _ := c.ClientSet.CoreV1().Pods(c.Namespace).Get(context.TODO(), c.Hostname, metav1.GetOptions{})
 	assert.Equal(t, map[string]string{"agent.mongodb.com/version": "5"}, thePod.Annotations)
-
-	os.Unsetenv(headlessAgent)
 }
 
 // TestHeadlessAgentReachedGoal verifies that the probe reports "true" if the config version is equal to the
 // last achieved version of the Agent
 func TestHeadlessAgentReachedGoal(t *testing.T) {
-	_ = os.Setenv(headlessAgent, "true")
+	t.Setenv(headlessAgent, "true")
 	c := testConfig("testdata/health-status-ok.json")
 	c.ClientSet = fake.NewSimpleClientset(testdata.TestPod(c.Namespace, c.Hostname), testdata.TestSecret(c.Namespace, c.AutomationConfigSecretName, 5))
 	ready, err := isPodReady(c)
 	assert.True(t, ready)
 	assert.NoError(t, err)
 	thePod, _ := c.ClientSet.CoreV1().Pods(c.Namespace).Get(context.TODO(), c.Hostname, metav1.GetOptions{})
 	assert.Equal(t, map[string]string{"agent.mongodb.com/version": "5"}, thePod.Annotations)
-
-	os.Unsetenv(headlessAgent)
 }
 
 func TestPodReadiness(t *testing.T) {

controllers/construct/build_statefulset_test.go

@@ -1,14 +1,14 @@
 package construct
 
 import (
-	"github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/podtemplatespec"
-	"github.com/mongodb/mongodb-kubernetes-operator/pkg/util/envvar"
 	"os"
 	"reflect"
 	"testing"
 
 	"github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/container"
+	"github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/podtemplatespec"
 	"github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/resourcerequirements"
+	"github.com/mongodb/mongodb-kubernetes-operator/pkg/util/envvar"
 
 	corev1 "k8s.io/api/core/v1"
 
@@ -40,9 +40,9 @@ func newTestReplicaSet() mdbv1.MongoDBCommunity {
 }
 
 func TestMultipleCalls_DoNotCauseSideEffects(t *testing.T) {
-	_ = os.Setenv(MongodbRepoUrl, "repo")
-	_ = os.Setenv(MongodbImageEnv, "mongo")
-	_ = os.Setenv(AgentImageEnv, "agent-image")
+	t.Setenv(MongodbRepoUrl, "repo")
+	t.Setenv(MongodbImageEnv, "mongo")
+	t.Setenv(AgentImageEnv, "agent-image")
 
 	mdb := newTestReplicaSet()
 	stsFunc := BuildMongoDBReplicaSetStatefulSetModificationFunction(&mdb, mdb)
@@ -63,10 +63,10 @@ func TestMultipleCalls_DoNotCauseSideEffects(t *testing.T) {
 }
 
 func TestManagedSecurityContext(t *testing.T) {
-	_ = os.Setenv(MongodbRepoUrl, "repo")
-	_ = os.Setenv(MongodbImageEnv, "mongo")
-	_ = os.Setenv(AgentImageEnv, "agent-image")
-	_ = os.Setenv(podtemplatespec.ManagedSecurityContextEnv, "true")
+	t.Setenv(MongodbRepoUrl, "repo")
+	t.Setenv(MongodbImageEnv, "mongo")
+	t.Setenv(AgentImageEnv, "agent-image")
+	t.Setenv(podtemplatespec.ManagedSecurityContextEnv, "true")
 
 	mdb := newTestReplicaSet()
 	stsFunc := BuildMongoDBReplicaSetStatefulSetModificationFunction(&mdb, mdb)

controllers/replicaset_controller_test.go

@@ -142,8 +142,8 @@ func TestKubernetesResources_AreCreated(t *testing.T) {
 }
 
 func TestStatefulSet_IsCorrectlyConfigured(t *testing.T) {
-	_ = os.Setenv(construct.MongodbRepoUrl, "repo")
-	_ = os.Setenv(construct.MongodbImageEnv, "mongo")
+	t.Setenv(construct.MongodbRepoUrl, "repo")
+	t.Setenv(construct.MongodbImageEnv, "mongo")
 
 	mdb := newTestReplicaSet()
 	mgr := client.NewManager(&mdb)

pkg/util/envvar/envvars_test.go

@@ -1,16 +1,13 @@
 package envvar
 
 import (
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestGetEnvOrDefault(t *testing.T) {
-
-	err := os.Setenv("env1", "val1")
-	assert.NoError(t, err)
+	t.Setenv("env1", "val1")
 
 	val := GetEnvOrDefault("env1", "defaultVal1")
 	assert.Equal(t, "val1", val)