-
Notifications
You must be signed in to change notification settings - Fork 683
/
tp.go
137 lines (120 loc) · 3.79 KB
/
tp.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
package testprocess
import (
"flag"
"log"
"os"
"os/exec"
"reflect"
"runtime"
"runtime/debug"
)
// flags. Initialize these in the init() step, in case anything else
// wants to call flag.Parse() from TestMain.
var (
name = flag.String("testprocess.name", "", "internal use")
)
// package-scoped global variables, that we use as regular variables
var (
functions = map[string]func(){}
dispatched = false
)
func getFunctionName(i interface{}) string {
return runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
}
func alreadySudoed() bool {
return os.Getuid() == 0 && os.Getenv("SUDO_USER") != ""
}
/* #nosec */
func _make(sudo bool, f func()) *exec.Cmd {
name := getFunctionName(f)
functions[name] = f
args := []string{os.Args[0], "-testprocess.name=" + name}
switch {
case sudo && !alreadySudoed():
args = append([]string{"sudo", "-E", "--"}, args...)
case !sudo && alreadySudoed():
// In case they called dtest.Sudo() to run "everything" as root.
args = append([]string{"sudo", "-E", "-u", os.Getenv("SUDO_USER"), "--"}, args...)
}
cmd := exec.Command(args[0], args[1:]...)
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
return cmd
}
// Dispatch can be used to launch multiple subprocesses as part of a
// go test. If you want to throw up in your mouth a little, then read
// the implementation. Don't worry, this is apparently a "blessed"
// hack for doing this sort of thing.
//
// You always want to call testprocess.Dispatch() at the beginning of
// your TestMain function since confusing things will happen if you
// call it later on or have complex logic surrounding calls to it. For
// each subprocess you want, use testprocess.Make to create an
// *exec.Cmd variable and save it in a global (it must be global for
// this to work). The resulting commands that are returned can be
// started/stopped at any point later in your test, e.g.:
//
// func TestMain(m *testing.M) {
// testprocess.Dispatch()
// os.Exit(m.Run())
// }
//
// var fooCmd = testprocess.Make(func() { doFoo(); })
// var barCmd = testprocess.Make(func() { doBar(); })
// var bazcmd = testprocess.Make(func() { doBaz(); })
//
// func TestSomething(t *testing.T) {
// ...
// err := fooCmd.Run()
// ...
// }
//
// It is permissible, but not required, to call flag.Parse before
// calling testprocess.Dispatch. If flag.Parse has not been called,
// then testprocess.Dispatch will call it.
func Dispatch() {
dispatched = true
if !flag.Parsed() {
flag.Parse()
}
if *name == "" {
return
}
log.Printf("TESTPROCESS %s PID: %d", *name, os.Getpid())
defer func() {
if r := recover(); r != nil {
stack := string(debug.Stack())
log.Printf("TESTPROCESS %s PANICKED: %v\n%s", *name, r, stack)
os.Exit(1)
}
}()
functions[*name]()
log.Printf("TESTPROCESS %s NORMAL EXIT", *name)
os.Exit(0)
}
// Make returns an *exec.Cmd that will execute the supplied function
// in a subprocess. For this to work, testprocess.Dispatch must be
// invoked by the TestMain of any test suite using this, and the call
// to Make must be *before* the call to testprocess.Dispatch; possibly
// from a global variable initializer, e.g.:
//
// var myCmd = testprocess.Make(func() { doSomething(); })
//
func Make(f func()) *exec.Cmd {
if dispatched {
// panic because it's a bug in the code, and a stack
// trace is useful.
panic("testprocess: testprocess.Make called after testprocess.Dispatch")
}
return _make(false, f)
}
// MakeSudo does the same thing as testprocess.Make with exactly the
// same limitations, except the subprocess will run as root.
func MakeSudo(f func()) *exec.Cmd {
if dispatched {
// panic because it's a bug in the code, and a stack
// trace is useful.
panic("testprocess: testprocess.MakeSudo called after testprocess.Dispatch")
}
return _make(true, f)
}