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