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  }