Add CLI flags to set environment variables

cmd/archive.go

@@ -34,12 +34,12 @@ var archiveCmd = &cobra.Command{
 	Use:   "archive",
 	Short: "Create an archive",
 	Long: `Create an archive.
-	 
+
 An archive is a fully self-contained test run, and can be executed identically elsewhere.`,
 	Example: `
   # Archive a test run.
   k6 archive -u 10 -d 10s -O myarchive.tar script.js
-  
+
   # Run the resulting archive.
   k6 run myarchive.tar`[1:],
 	Args: cobra.ExactArgs(1),
@@ -55,7 +55,13 @@ An archive is a fully self-contained test run, and can be executed identically e
 		if err != nil {
 			return err
 		}
-		r, err := newRunner(src, runType, afero.NewOsFs())
+
+		runtimeOptions, err := getRuntimeOptions(cmd.Flags())
+		if err != nil {
+			return err
+		}
+
+		r, err := newRunner(src, runType, afero.NewOsFs(), runtimeOptions)
 		if err != nil {
 			return err
 		}
@@ -88,6 +94,8 @@ An archive is a fully self-contained test run, and can be executed identically e
 
 func init() {
 	RootCmd.AddCommand(archiveCmd)
+	archiveCmd.Flags().SortFlags = false
 	archiveCmd.Flags().AddFlagSet(optionFlagSet())
+	archiveCmd.Flags().AddFlagSet(runtimeOptionFlagSet(false))
 	archiveCmd.Flags().StringVarP(&archiveOut, "archive-out", "O", archiveOut, "archive output filename")
 }

cmd/cloud.go

@@ -60,7 +60,12 @@ This will execute the test on the Load Impact cloud service. Use "k6 login cloud
 			return err
 		}
 
-		r, err := newRunner(src, runType, afero.NewOsFs())
+		runtimeOptions, err := getRuntimeOptions(cmd.Flags())
+		if err != nil {
+			return err
+		}
+
+		r, err := newRunner(src, runType, afero.NewOsFs(), runtimeOptions)
 		if err != nil {
 			return err
 		}
@@ -136,5 +141,7 @@ This will execute the test on the Load Impact cloud service. Use "k6 login cloud
 
 func init() {
 	RootCmd.AddCommand(cloudCmd)
+	cloudCmd.Flags().SortFlags = false
 	cloudCmd.Flags().AddFlagSet(optionFlagSet())
+	cloudCmd.Flags().AddFlagSet(runtimeOptionFlagSet(false))
 }

cmd/inspect.go

@@ -54,16 +54,25 @@ var inspectCmd = &cobra.Command{
 			typ = detectType(src.Data)
 		}
 
+		runtimeOptions, err := getRuntimeOptions(cmd.Flags())
+		if err != nil {
+			return err
+		}
+
 		var opts lib.Options
 		switch typ {
 		case typeArchive:
 			arc, err := lib.ReadArchive(bytes.NewBuffer(src.Data))
 			if err != nil {
 				return err
 			}
-			opts = arc.Options
+			b, err := js.NewBundleFromArchive(arc, runtimeOptions)
+			if err != nil {
+				return err
+			}
+			opts = b.Options
 		case typeJS:
-			b, err := js.NewBundle(src, fs)
+			b, err := js.NewBundle(src, fs, runtimeOptions)
 			if err != nil {
 				return err
 			}
@@ -81,5 +90,7 @@ var inspectCmd = &cobra.Command{
 
 func init() {
 	RootCmd.AddCommand(inspectCmd)
+	inspectCmd.Flags().SortFlags = false
+	inspectCmd.Flags().AddFlagSet(runtimeOptionFlagSet(false))
 	inspectCmd.Flags().StringVarP(&runType, "type", "t", runType, "override file `type`, \"js\" or \"archive\"")
 }

cmd/run.go

@@ -105,7 +105,13 @@ a commandline interface for interacting with it.`,
 		if err != nil {
 			return err
 		}
-		r, err := newRunner(src, runType, afero.NewOsFs())
+
+		runtimeOptions, err := getRuntimeOptions(cmd.Flags())
+		if err != nil {
+			return err
+		}
+
+		r, err := newRunner(src, runType, afero.NewOsFs(), runtimeOptions)
 		if err != nil {
 			return err
 		}
@@ -408,6 +414,7 @@ func init() {
 
 	runCmd.Flags().SortFlags = false
 	runCmd.Flags().AddFlagSet(optionFlagSet())
+	runCmd.Flags().AddFlagSet(runtimeOptionFlagSet(true))
 	runCmd.Flags().AddFlagSet(configFlagSet())
 	runCmd.Flags().StringVarP(&runType, "type", "t", runType, "override file `type`, \"js\" or \"archive\"")
 }
@@ -429,20 +436,20 @@ func readSource(src, pwd string, fs afero.Fs, stdin io.Reader) (*lib.SourceData,
 }
 
 // Creates a new runner.
-func newRunner(src *lib.SourceData, typ string, fs afero.Fs) (lib.Runner, error) {
+func newRunner(src *lib.SourceData, typ string, fs afero.Fs, rtOpts lib.RuntimeOptions) (lib.Runner, error) {
 	switch typ {
 	case "":
-		return newRunner(src, detectType(src.Data), fs)
+		return newRunner(src, detectType(src.Data), fs, rtOpts)
 	case typeJS:
-		return js.New(src, fs)
+		return js.New(src, fs, rtOpts)
 	case typeArchive:
 		arc, err := lib.ReadArchive(bytes.NewReader(src.Data))
 		if err != nil {
 			return nil, err
 		}
 		switch arc.Type {
 		case typeJS:
-			return js.NewFromArchive(arc)
+			return js.NewFromArchive(arc, rtOpts)
 		default:
 			return nil, errors.Errorf("archive requests unsupported runner: %s", arc.Type)
 		}

cmd/runtime_options.go

@@ -0,0 +1,87 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2018 Load Impact
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+package cmd
+
+import (
+	"os"
+	"regexp"
+	"strings"
+
+	"github.com/loadimpact/k6/lib"
+	"github.com/pkg/errors"
+	"github.com/spf13/pflag"
+)
+
+var userEnvVarName = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)
+
+func parseEnvKeyValue(kv string) (string, string) {
+	if idx := strings.IndexRune(kv, '='); idx != -1 {
+		return kv[:idx], kv[idx+1:]
+	}
+	return kv, ""
+}
+
+func collectEnv() map[string]string {
+	env := make(map[string]string)
+	for _, kv := range os.Environ() {
+		k, v := parseEnvKeyValue(kv)
+		env[k] = v
+	}
+	return env
+}
+
+func runtimeOptionFlagSet(includeSysEnv bool) *pflag.FlagSet {
+	flags := pflag.NewFlagSet("", 0)
+	flags.SortFlags = false
+	flags.Bool("include-system-env-vars", includeSysEnv, "pass the real system environment variables to the runtime")
+	flags.StringSliceP("env", "e", nil, "add/override environment variable with `VAR=value`")
+	return flags
+}
+
+func getRuntimeOptions(flags *pflag.FlagSet) (lib.RuntimeOptions, error) {
+	opts := lib.RuntimeOptions{
+		IncludeSystemEnvVars: getNullBool(flags, "include-system-env-vars"),
+		Env:                  make(map[string]string),
+	}
+
+	// If enabled, gather the actual system environment variables
+	if opts.IncludeSystemEnvVars.Bool {
+		opts.Env = collectEnv()
+	}
+
+	// Set/overwrite environment variables with custom user-supplied values
+	envVars, err := flags.GetStringSlice("env")
+	if err != nil {
+		return opts, err
+	}
+	if len(envVars) > 0 {
+		for _, kv := range envVars {
+			k, v := parseEnvKeyValue(kv)
+			// Allow only alphanumeric ASCII variable names for now
+			if !userEnvVarName.MatchString(k) {
+				return opts, errors.Errorf("Invalid environment variable name '%s'", k)
+			}
+			opts.Env[k] = v
+		}
+	}
+
+	return opts, nil
+}