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 }