test(memlimit): rollback environment variable after testing

KimMachineGun · Jan 9, 2024 · e243cf4 · e243cf4
1 parent ab8c940
commit e243cf4
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/memlimit/experiment_test.go b/memlimit/experiment_test.go
@@ -45,6 +45,15 @@ func TestParseExperiments(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			exp, ok := os.LookupEnv(envAUTOMEMLIMIT_EXPERIMENT)
+			t.Cleanup(func() {
+				if ok {
+					os.Setenv(envAUTOMEMLIMIT_EXPERIMENT, exp)
+				} else {
+					os.Unsetenv(envAUTOMEMLIMIT_EXPERIMENT)
+				}
+			})
+
 			os.Setenv("AUTOMEMLIMIT_EXPERIMENT", tt.env)
 			exps, err := parseExperiments()
 			if !reflect.DeepEqual(exps, tt.want) {