github.com/jrxfive/nomad@v0.6.1-0.20170802162750-1fef470e89bf/helper/testtask/testtask.go (about)

     1  // Package testtask implements a portable set of commands useful as stand-ins
     2  // for user tasks.
     3  package testtask
     4  
     5  import (
     6  	"fmt"
     7  	"io/ioutil"
     8  	"os"
     9  	"os/exec"
    10  	"time"
    11  
    12  	"github.com/hashicorp/nomad/nomad/structs"
    13  	"github.com/kardianos/osext"
    14  )
    15  
    16  // Path returns the path to the currently running executable.
    17  func Path() string {
    18  	path, err := osext.Executable()
    19  	if err != nil {
    20  		panic(err)
    21  	}
    22  	return path
    23  }
    24  
    25  // SetCmdEnv configures the environment of cmd so that Run executes a testtask
    26  // script when called from within cmd.
    27  func SetCmdEnv(cmd *exec.Cmd) {
    28  	cmd.Env = append(os.Environ(), "TEST_TASK=execute")
    29  }
    30  
    31  // SetTaskEnv configures the environment of t so that Run executes a testtask
    32  // script when called from within t.
    33  func SetTaskEnv(t *structs.Task) {
    34  	if t.Env == nil {
    35  		t.Env = map[string]string{}
    36  	}
    37  	t.Env["TEST_TASK"] = "execute"
    38  }
    39  
    40  // Run interprets os.Args as a testtask script if the current program was
    41  // launched with an environment configured by SetCmdEnv or SetTaskEnv. It
    42  // returns false if the environment was not set by this package.
    43  func Run() bool {
    44  	switch tm := os.Getenv("TEST_TASK"); tm {
    45  	case "":
    46  		return false
    47  	case "execute":
    48  		execute()
    49  		return true
    50  	default:
    51  		fmt.Fprintf(os.Stderr, "unexpected value for TEST_TASK, \"%s\"\n", tm)
    52  		os.Exit(1)
    53  		return true
    54  	}
    55  }
    56  
    57  func execute() {
    58  	if len(os.Args) < 2 {
    59  		fmt.Fprintln(os.Stderr, "no command provided")
    60  		os.Exit(1)
    61  	}
    62  
    63  	args := os.Args[1:]
    64  
    65  	// popArg removes the first argument from args and returns it.
    66  	popArg := func() string {
    67  		s := args[0]
    68  		args = args[1:]
    69  		return s
    70  	}
    71  
    72  	// execute a sequence of operations from args
    73  	for len(args) > 0 {
    74  		switch cmd := popArg(); cmd {
    75  
    76  		case "sleep":
    77  			// sleep <dur>: sleep for a duration indicated by the first
    78  			// argument
    79  			if len(args) < 1 {
    80  				fmt.Fprintln(os.Stderr, "expected arg for sleep")
    81  				os.Exit(1)
    82  			}
    83  			dur, err := time.ParseDuration(popArg())
    84  			if err != nil {
    85  				fmt.Fprintf(os.Stderr, "could not parse sleep time: %v", err)
    86  				os.Exit(1)
    87  			}
    88  			time.Sleep(dur)
    89  
    90  		case "echo":
    91  			// echo <msg>: write the msg followed by a newline to stdout.
    92  			fmt.Println(popArg())
    93  
    94  		case "write":
    95  			// write <msg> <file>: write a message to a file. The first
    96  			// argument is the msg. The second argument is the path to the
    97  			// target file.
    98  			if len(args) < 2 {
    99  				fmt.Fprintln(os.Stderr, "expected two args for write")
   100  				os.Exit(1)
   101  			}
   102  			msg := popArg()
   103  			file := popArg()
   104  			ioutil.WriteFile(file, []byte(msg), 0666)
   105  
   106  		default:
   107  			fmt.Fprintln(os.Stderr, "unknown command:", cmd)
   108  			os.Exit(1)
   109  		}
   110  	}
   111  }