added support for KOI_NS and KOI_CTX env vars to define context and n…

…amespace

koi/tweakArgs.go

@@ -1,9 +1,16 @@
 package koi
 
 import (
+	"os"
 	"strings"
 )
 
+type defaultValueMapping struct {
+	value          string
+	alreadySet     bool
+	flagsThatMatch []string
+}
+
 // applyTweaksToArgs modifies the args sent in so they work wtih kubectl
 // Goals:
 // The -x flag should become --context
@@ -13,18 +20,76 @@ func ApplyTweaksToArgs(args []string) []string {
 		"-x": "--context",
 	}
 
+	defaultValuesForFlags := []*defaultValueMapping{
+		{
+			value: coalesceString(os.Getenv("KOI_CONTEXT"), os.Getenv("KOI_CTX")),
+			flagsThatMatch: []string{
+				"--context",
+				"-x",
+			},
+		},
+		{
+			value: coalesceString(os.Getenv("KOI_NAMESPACE"), os.Getenv("KOI_NS")),
+			flagsThatMatch: []string{
+				"--namespace",
+				"-n",
+			},
+		},
+	}
+
+	lastArgPos := 0
+
 	for i, arg := range args {
+		lastArgPos = i
 		if arg == "--" {
 			break
 		}
+
+		// If any of the args match teh shorthand, then we set it to the longform
 		for shorthand, longform := range shortHandReplacements {
 			if strings.HasPrefix(arg, shorthand+"=") || arg == shorthand {
 				if _, ok := replacedFlags[shorthand]; !ok {
 					args[i] = strings.Replace(arg, shorthand, longform, 1)
+					arg = args[i]
 					replacedFlags[shorthand] = true
 				}
 			}
 		}
+		// If any of the args match one of the defautl values, then we don't need to apply them
+		for _, dv := range defaultValuesForFlags {
+			if !dv.alreadySet && dv.value != "" && stringArrayContains(dv.flagsThatMatch, arg) {
+				dv.alreadySet = true
+			}
+		}
+	}
+
+	outputArgs := make([]string, 0, len(args)+len(defaultValuesForFlags))
+
+	outputArgs = append(outputArgs, args[0:lastArgPos]...)
+	for _, dv := range defaultValuesForFlags {
+		if dv.value != "" && !dv.alreadySet {
+			outputArgs = append(outputArgs, dv.flagsThatMatch[0], dv.value)
+		}
+	}
+	outputArgs = append(outputArgs, args[lastArgPos:]...)
+
+	return outputArgs
+}
+
+func stringArrayContains(haystack []string, needle string) bool {
+	for _, s := range haystack {
+		if s == needle {
+			return true
+		}
+	}
+	return false
+}
+
+func coalesceString(check ...string) string {
+	for _, v := range check {
+		if v != "" {
+			return v
+		}
 	}
-	return args
+	return ""
 }

koi/tweakArgs_test.go

@@ -1,6 +1,7 @@
 package koi
 
 import (
+	"os"
 	"reflect"
 	"testing"
 )
@@ -10,6 +11,7 @@ func Test_ApplyTweaksToArgs(t *testing.T) {
 		name string
 		args []string
 		want []string
+		env  map[string]string
 	}{
 		{
 			name: "-x flag should be changed to --context",
@@ -31,11 +33,58 @@ func Test_ApplyTweaksToArgs(t *testing.T) {
 			args: []string{"-n", "bob", "--", "bash", "-x", "example"},
 			want: []string{"-n", "bob", "--", "bash", "-x", "example"},
 		},
+		{
+			name: "If KOI_NS is set, then the namespace should be set to that if it is not alreayd set",
+			args: []string{},
+			want: []string{"--namespace", "koi"},
+			env: map[string]string{
+				"KOI_NS": "koi",
+			},
+		},
+		{
+			name: "If KOI_CONTEXT is set, then the context should be set to that if it is not already set",
+			args: []string{},
+			want: []string{"--context", "koi"},
+			env: map[string]string{
+				"KOI_CONTEXT": "koi",
+			},
+		},
+		{
+			name: "If KOI_CONTEXT is set, then the context should be set to that if it is not already set",
+			args: []string{"--context", "bob"},
+			want: []string{"--context", "bob"},
+			env: map[string]string{
+				"KOI_CONTEXT": "koi",
+			},
+		},
+		{
+			name: "If KOI_CONTEXT is set, then the context should be set before double-dash",
+			args: []string{"exec", "-it", "--", "--context", "ignroethis"},
+			want: []string{"exec", "-it", "--context", "koi", "--", "--context", "ignroethis"},
+			env: map[string]string{
+				"KOI_CONTEXT": "koi",
+			},
+		},
+		{
+			name: "If both KOI_CONTEXt and KOI_NS are set, then they should both be used",
+			args: []string{"exec", "-it", "--", "--context", "ignroethis"},
+			want: []string{"exec", "-it", "--context", "koi", "--namespace", "koi", "--", "--context", "ignroethis"},
+			env: map[string]string{
+				"KOI_CTX": "koi",
+				"KOI_NS":  "koi",
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			for key, val := range tt.env {
+				os.Setenv(key, val)
+			}
 			if got := ApplyTweaksToArgs(tt.args); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("applyTweaksToArgs() = %v, want %v", got, tt.want)
+				t.Errorf("applyTweaksToArgs() got: %v, want: %v", got, tt.want)
+			}
+			for key := range tt.env {
+				os.Unsetenv(key)
 			}
 		})
 	}