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