...

Source file src/github.com/rogpeppe/go-internal/testscript/exe.go

Documentation: github.com/rogpeppe/go-internal/testscript

     1  // Copyright 2018 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package testscript
     6  
     7  import (
     8  	"io"
     9  	"io/ioutil"
    10  	"log"
    11  	"os"
    12  	"path/filepath"
    13  	"runtime"
    14  	"strings"
    15  )
    16  
    17  // TestingM is implemented by *testing.M. It's defined as an interface
    18  // to allow testscript to co-exist with other testing frameworks
    19  // that might also wish to call M.Run.
    20  type TestingM interface {
    21  	Run() int
    22  }
    23  
    24  // Deprecated: this option is no longer used.
    25  func IgnoreMissedCoverage() {}
    26  
    27  // RunMain should be called within a TestMain function to allow
    28  // subcommands to be run in the testscript context.
    29  //
    30  // The commands map holds the set of command names, each
    31  // with an associated run function which should return the
    32  // code to pass to os.Exit. It's OK for a command function to
    33  // exit itself, but this may result in loss of coverage information.
    34  //
    35  // When Run is called, these commands are installed as regular commands in the shell
    36  // path, so can be invoked with "exec" or via any other command (for example a shell script).
    37  //
    38  // For backwards compatibility, the commands declared in the map can be run
    39  // without "exec" - that is, "foo" will behave like "exec foo".
    40  // This can be disabled with Params.RequireExplicitExec to keep consistency
    41  // across test scripts, and to keep separate process executions explicit.
    42  //
    43  // This function returns an exit code to pass to os.Exit, after calling m.Run.
    44  func RunMain(m TestingM, commands map[string]func() int) (exitCode int) {
    45  	// Depending on os.Args[0], this is either the top-level execution of
    46  	// the test binary by "go test", or the execution of one of the provided
    47  	// commands via "foo" or "exec foo".
    48  
    49  	cmdName := filepath.Base(os.Args[0])
    50  	if runtime.GOOS == "windows" {
    51  		cmdName = strings.TrimSuffix(cmdName, ".exe")
    52  	}
    53  	mainf := commands[cmdName]
    54  	if mainf == nil {
    55  		// Unknown command; this is just the top-level execution of the
    56  		// test binary by "go test".
    57  
    58  		// Set up all commands in a directory, added in $PATH.
    59  		tmpdir, err := ioutil.TempDir("", "testscript-main")
    60  		if err != nil {
    61  			log.Printf("could not set up temporary directory: %v", err)
    62  			return 2
    63  		}
    64  		defer func() {
    65  			if err := os.RemoveAll(tmpdir); err != nil {
    66  				log.Printf("cannot delete temporary directory: %v", err)
    67  				exitCode = 2
    68  			}
    69  		}()
    70  		bindir := filepath.Join(tmpdir, "bin")
    71  		if err := os.MkdirAll(bindir, 0o777); err != nil {
    72  			log.Printf("could not set up PATH binary directory: %v", err)
    73  			return 2
    74  		}
    75  		os.Setenv("PATH", bindir+string(filepath.ListSeparator)+os.Getenv("PATH"))
    76  
    77  		// We're not in a subcommand.
    78  		for name := range commands {
    79  			name := name
    80  			// Set up this command in the directory we added to $PATH.
    81  			binfile := filepath.Join(bindir, name)
    82  			if runtime.GOOS == "windows" {
    83  				binfile += ".exe"
    84  			}
    85  			binpath, err := os.Executable()
    86  			if err == nil {
    87  				err = copyBinary(binpath, binfile)
    88  			}
    89  			if err != nil {
    90  				log.Printf("could not set up %s in $PATH: %v", name, err)
    91  				return 2
    92  			}
    93  			scriptCmds[name] = func(ts *TestScript, neg bool, args []string) {
    94  				if ts.params.RequireExplicitExec {
    95  					ts.Fatalf("use 'exec %s' rather than '%s' (because RequireExplicitExec is enabled)", name, name)
    96  				}
    97  				ts.cmdExec(neg, append([]string{name}, args...))
    98  			}
    99  		}
   100  		return m.Run()
   101  	}
   102  	// The command being registered is being invoked, so run it, then exit.
   103  	os.Args[0] = cmdName
   104  	return mainf()
   105  }
   106  
   107  // copyBinary makes a copy of a binary to a new location. It is used as part of
   108  // setting up top-level commands in $PATH.
   109  //
   110  // It does not attempt to use symlinks for two reasons:
   111  //
   112  // First, some tools like cmd/go's -toolexec will be clever enough to realise
   113  // when they're given a symlink, and they will use the symlink target for
   114  // executing the program. This breaks testscript, as we depend on os.Args[0] to
   115  // know what command to run.
   116  //
   117  // Second, symlinks might not be available on some environments, so we have to
   118  // implement a "full copy" fallback anyway.
   119  //
   120  // However, we do try to use cloneFile, since that will probably work on most
   121  // unix-like setups. Note that "go test" also places test binaries in the
   122  // system's temporary directory, like we do.
   123  func copyBinary(from, to string) error {
   124  	if err := cloneFile(from, to); err == nil {
   125  		return nil
   126  	}
   127  	writer, err := os.OpenFile(to, os.O_WRONLY|os.O_CREATE, 0o777)
   128  	if err != nil {
   129  		return err
   130  	}
   131  	defer writer.Close()
   132  
   133  	reader, err := os.Open(from)
   134  	if err != nil {
   135  		return err
   136  	}
   137  	defer reader.Close()
   138  
   139  	_, err = io.Copy(writer, reader)
   140  	return err
   141  }
   142  

View as plain text