Skip to content

Commit 37894e9

Browse files
JuneezeeJuneezee
committed
test: use T.Setenv to set env vars in tests
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 <[email protected]>
1 parent 835767a commit 37894e9

File tree

6 files changed

+13
-57
lines changed

6 files changed

+13
-57
lines changed

‎cpu/cpu_linux_test.go

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,14 @@ package cpu
22

33
import (
44
"errors"
5-
"os"
65
"os/exec"
76
"strconv"
87
"strings"
98
"testing"
109
)
1110

1211
func TestTimesEmpty(t *testing.T) {
13-
orig := os.Getenv("HOST_PROC")
14-
os.Setenv("HOST_PROC", "testdata/linux/times_empty")
12+
t.Setenv("HOST_PROC", "testdata/linux/times_empty")
1513
_, err := Times(true)
1614
if err != nil {
1715
t.Error("Times(true) failed")
@@ -20,12 +18,10 @@ func TestTimesEmpty(t *testing.T) {
2018
if err != nil {
2119
t.Error("Times(false) failed")
2220
}
23-
os.Setenv("HOST_PROC", orig)
2421
}
2522

2623
func TestCPUparseStatLine_424(t *testing.T) {
27-
orig := os.Getenv("HOST_PROC")
28-
os.Setenv("HOST_PROC", "testdata/linux/424/proc")
24+
t.Setenv("HOST_PROC", "testdata/linux/424/proc")
2925
{
3026
l, err := Times(true)
3127
if err != nil || len(l) == 0 {
@@ -40,7 +36,6 @@ func TestCPUparseStatLine_424(t *testing.T) {
4036
}
4137
t.Logf("Times(false): %#v", l)
4238
}
43-
os.Setenv("HOST_PROC", orig)
4439
}
4540

4641
func TestCPUCountsAgainstLscpu(t *testing.T) {
@@ -93,9 +88,7 @@ func TestCPUCountsAgainstLscpu(t *testing.T) {
9388
}
9489

9590
func TestCPUCountsLogicalAndroid_1037(t *testing.T) { // https://github.com/shirou/gopsutil/issues/1037
96-
orig := os.Getenv("HOST_PROC")
97-
os.Setenv("HOST_PROC", "testdata/linux/1037/proc")
98-
defer os.Setenv("HOST_PROC", orig)
91+
t.Setenv("HOST_PROC", "testdata/linux/1037/proc")
9992

10093
count, err := Counts(true)
10194
if err != nil {

‎cpu/cpu_plan9_test.go

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
package cpu
55

66
import (
7-
"os"
87
"path/filepath"
98
"testing"
109

@@ -30,13 +29,9 @@ var timesTests = []struct {
3029
}
3130

3231
func TestTimesPlan9(t *testing.T) {
33-
origRoot := os.Getenv("HOST_ROOT")
34-
t.Cleanup(func() {
35-
os.Setenv("HOST_ROOT", origRoot)
36-
})
3732
for _, tt := range timesTests {
3833
t.Run(tt.mockedRootFS, func(t *testing.T) {
39-
os.Setenv("HOST_ROOT", filepath.Join("testdata/plan9", tt.mockedRootFS))
34+
t.Setenv("HOST_ROOT", filepath.Join("testdata/plan9", tt.mockedRootFS))
4035
stats, err := Times(false)
4136
skipIfNotImplementedErr(t, err)
4237
if err != nil {

‎internal/common/common.go

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -364,16 +364,6 @@ func HostDev(combineWith ...string) string {
364364
return GetEnv("HOST_DEV", "/dev", combineWith...)
365365
}
366366

367-
// MockEnv set environment variable and return revert function.
368-
// MockEnv should be used testing only.
369-
func MockEnv(key string, value string) func() {
370-
original := os.Getenv(key)
371-
os.Setenv(key, value)
372-
return func() {
373-
os.Setenv(key, original)
374-
}
375-
}
376-
377367
// getSysctrlEnv sets LC_ALL=C in a list of env vars for use when running
378368
// sysctl commands (see DoSysctrl).
379369
func getSysctrlEnv(env []string) []string {

‎mem/mem_linux_test.go

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
package mem
55

66
import (
7-
"os"
87
"path/filepath"
98
"reflect"
109
"strings"
@@ -111,12 +110,9 @@ var virtualMemoryTests = []struct {
111110
}
112111

113112
func TestVirtualMemoryLinux(t *testing.T) {
114-
origProc := os.Getenv("HOST_PROC")
115-
defer os.Setenv("HOST_PROC", origProc)
116-
117113
for _, tt := range virtualMemoryTests {
118114
t.Run(tt.mockedRootFS, func(t *testing.T) {
119-
os.Setenv("HOST_PROC", filepath.Join("testdata/linux/virtualmemory/", tt.mockedRootFS, "proc"))
115+
t.Setenv("HOST_PROC", filepath.Join("testdata/linux/virtualmemory/", tt.mockedRootFS, "proc"))
120116

121117
stat, err := VirtualMemory()
122118
skipIfNotImplementedErr(t, err)

‎mem/mem_plan9_test.go

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
package mem
55

66
import (
7-
"os"
87
"reflect"
98
"testing"
109
)
@@ -27,14 +26,9 @@ var virtualMemoryTests = []struct {
2726
}
2827

2928
func TestVirtualMemoryPlan9(t *testing.T) {
30-
origProc := os.Getenv("HOST_ROOT")
31-
t.Cleanup(func() {
32-
os.Setenv("HOST_ROOT", origProc)
33-
})
34-
3529
for _, tt := range virtualMemoryTests {
3630
t.Run(tt.mockedRootFS, func(t *testing.T) {
37-
os.Setenv("HOST_ROOT", "testdata/plan9/virtualmemory/")
31+
t.Setenv("HOST_ROOT", "testdata/plan9/virtualmemory/")
3832

3933
stat, err := VirtualMemory()
4034
skipIfNotImplementedErr(t, err)
@@ -62,14 +56,9 @@ var swapMemoryTests = []struct {
6256
}
6357

6458
func TestSwapMemoryPlan9(t *testing.T) {
65-
origProc := os.Getenv("HOST_ROOT")
66-
t.Cleanup(func() {
67-
os.Setenv("HOST_ROOT", origProc)
68-
})
69-
7059
for _, tt := range swapMemoryTests {
7160
t.Run(tt.mockedRootFS, func(t *testing.T) {
72-
os.Setenv("HOST_ROOT", "testdata/plan9/virtualmemory/")
61+
t.Setenv("HOST_ROOT", "testdata/plan9/virtualmemory/")
7362

7463
swap, err := SwapMemory()
7564
skipIfNotImplementedErr(t, err)

‎process/process_linux_test.go

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import (
1111
"strings"
1212
"testing"
1313

14-
"github.com/shirou/gopsutil/v3/internal/common"
1514
"github.com/stretchr/testify/assert"
1615
)
1716

@@ -62,8 +61,7 @@ func Test_Process_splitProcStat_fromFile(t *testing.T) {
6261
if err != nil {
6362
t.Error(err)
6463
}
65-
f := common.MockEnv("HOST_PROC", "testdata/linux")
66-
defer f()
64+
t.Setenv("HOST_PROC", "testdata/linux")
6765
for _, pid := range pids {
6866
pid, err := strconv.ParseInt(pid.Name(), 0, 32)
6967
if err != nil {
@@ -99,8 +97,7 @@ func Test_fillFromCommWithContext(t *testing.T) {
9997
if err != nil {
10098
t.Error(err)
10199
}
102-
f := common.MockEnv("HOST_PROC", "testdata/linux")
103-
defer f()
100+
t.Setenv("HOST_PROC", "testdata/linux")
104101
for _, pid := range pids {
105102
pid, err := strconv.ParseInt(pid.Name(), 0, 32)
106103
if err != nil {
@@ -121,8 +118,7 @@ func Test_fillFromStatusWithContext(t *testing.T) {
121118
if err != nil {
122119
t.Error(err)
123120
}
124-
f := common.MockEnv("HOST_PROC", "testdata/linux")
125-
defer f()
121+
t.Setenv("HOST_PROC", "testdata/linux")
126122
for _, pid := range pids {
127123
pid, err := strconv.ParseInt(pid.Name(), 0, 32)
128124
if err != nil {
@@ -139,8 +135,7 @@ func Test_fillFromStatusWithContext(t *testing.T) {
139135
}
140136

141137
func Benchmark_fillFromCommWithContext(b *testing.B) {
142-
f := common.MockEnv("HOST_PROC", "testdata/linux")
143-
defer f()
138+
b.Setenv("HOST_PROC", "testdata/linux")
144139
pid := 1060
145140
p, _ := NewProcess(int32(pid))
146141
for i := 0; i < b.N; i++ {
@@ -149,8 +144,7 @@ func Benchmark_fillFromCommWithContext(b *testing.B) {
149144
}
150145

151146
func Benchmark_fillFromStatusWithContext(b *testing.B) {
152-
f := common.MockEnv("HOST_PROC", "testdata/linux")
153-
defer f()
147+
b.Setenv("HOST_PROC", "testdata/linux")
154148
pid := 1060
155149
p, _ := NewProcess(int32(pid))
156150
for i := 0; i < b.N; i++ {
@@ -163,8 +157,7 @@ func Test_fillFromTIDStatWithContext_lx_brandz(t *testing.T) {
163157
if err != nil {
164158
t.Error(err)
165159
}
166-
f := common.MockEnv("HOST_PROC", "testdata/lx_brandz")
167-
defer f()
160+
t.Setenv("HOST_PROC", "testdata/lx_brandz")
168161
for _, pid := range pids {
169162
pid, err := strconv.ParseInt(pid.Name(), 0, 32)
170163
if err != nil {

0 commit comments

Comments
 (0)