...

Source file src/github.com/datawire/ambassador/v2/pkg/dtest/testprocess/tp.go

Documentation: github.com/datawire/ambassador/v2/pkg/dtest/testprocess

     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  

View as plain text