Correctly parses all arguments after kubectl/ssh and adds tests

Signed-off-by: JoshVanL <vleeuwenjoshua@gmail.com>

cmd/tarmak/cmd/cluster_instances_ssh.go

@@ -8,13 +8,14 @@ import (
 )
 
 var clusterInstancesSshCmd = &cobra.Command{
-	Use:   "ssh [instance alias]",
+	Use:   "ssh [instance alias] [optional ssh command]",
 	Short: "Log into an instance with SSH",
 	Run: func(cmd *cobra.Command, args []string) {
 		t := tarmak.New(globalFlags)
 		defer t.Cleanup()
 		t.SSHPassThrough(args)
 	},
+	DisableFlagsInUseLine: true,
 }
 
 func init() {

cmd/tarmak/cmd/cluster_kubectl.go

@@ -14,9 +14,9 @@ var clusterKubectlCmd = &cobra.Command{
 		t := tarmak.New(globalFlags)
 		t.Perform(t.NewCmdTarmak(args).Kubectl())
 	},
+	DisableFlagsInUseLine: true,
 }
 
 func init() {
-	clusterKubectlCmd.Flags().SetInterspersed(false)
 	clusterCmd.AddCommand(clusterKubectlCmd)
 }

cmd/tarmak/cmd/cluster_ssh.go

@@ -8,13 +8,14 @@ import (
 )
 
 var clusterSshCmd = &cobra.Command{
-	Use:   "ssh [instance alias]",
+	Use:   "ssh [instance alias] [optional ssh command]",
 	Short: "Log into an instance with SSH",
 	Run: func(cmd *cobra.Command, args []string) {
 		t := tarmak.New(globalFlags)
 		defer t.Cleanup()
 		t.SSHPassThrough(args)
 	},
+	DisableFlagsInUseLine: true,
 }
 
 func init() {

cmd/tarmak/cmd/root.go

@@ -10,6 +10,7 @@ import (
 	"github.com/spf13/cobra"
 
 	tarmakv1alpha1 "github.com/jetstack/tarmak/pkg/apis/tarmak/v1alpha1"
+	"github.com/jetstack/tarmak/pkg/tarmak/utils"
 	"github.com/jetstack/tarmak/pkg/terraform"
 )
 
@@ -38,20 +39,31 @@ func Execute(args []string) {
 		}
 	}
 
-	RootCmd.SetArgs(args)
+	cmd, _, err := RootCmd.Traverse(args)
+	if err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+	}
 
-	// evalutate command that is gonna be run
-	command, commandArgs, err := RootCmd.Traverse(args)
+	for _, c := range []struct {
+		use  string
+		name string
+	}{
+		{use: clusterKubectlCmd.Use, name: "kubectl"},
+		{use: clusterSshCmd.Use, name: "ssh"},
+	} {
+		if cmd.Use != c.use {
+			continue
+		}
 
-	// escape pass through commands (kubectl, ssh) if necessary
-	if err == nil && (command.Use == clusterKubectlCmd.Use || command.Use == clusterSshCmd.Use) {
-		// if no escape exists already add one
-		if !stringSliceContains(commandArgs, "--") {
-			pos := len(args) - len(commandArgs)
-			newArgs := append(args[:pos], append([]string{"--"}, args[pos:]...)...)
-			RootCmd.SetArgs(newArgs)
-			// this line helps debugging fmt.Printf("rewriting args\noriginal args=%v\ncommand  args=%v\nnew      args=%v\n", args, commandArgs, newArgs)
+		i := utils.IndexOfString(args, c.name)
+		if i == -1 {
+			break
 		}
+
+		RootCmd.SetArgs(
+			append(args[:i+1], append([]string{"--"}, args[i+1:]...)...))
+		break
 	}
 
 	if err := RootCmd.Execute(); err != nil {
@@ -60,15 +72,6 @@ func Execute(args []string) {
 	}
 }
 
-func stringSliceContains(s []string, e string) bool {
-	for _, a := range s {
-		if a == e {
-			return true
-		}
-	}
-	return false
-}
-
 func init() {
 	// set default tarmak config folder
 	tarmakConfigPath := "~/.tarmak"

cmd/tarmak/cmd/tarmak_test.go

@@ -0,0 +1,106 @@
+// Copyright Jetstack Ltd. See LICENSE for details.
+package cmd
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/spf13/cobra"
+)
+
+const (
+	globalFlag = "--current-cluster=xxx"
+)
+
+type cmdTest struct {
+	*testing.T
+	cmd    *cobra.Command
+	argsCh chan string
+	args   []string
+}
+
+func newCmdTest(t *testing.T, cmd *cobra.Command) *cmdTest {
+	return &cmdTest{
+		T:   t,
+		cmd: cmd,
+	}
+}
+
+func (c *cmdTest) testRunFunc() func(cmd *cobra.Command, args []string) {
+	return func(cmd *cobra.Command, args []string) {
+		for _, a := range args {
+			c.argsCh <- a
+		}
+		close(c.argsCh)
+	}
+}
+
+func (c *cmdTest) flagIgnored(exp bool) {
+	c.argsCh = make(chan string)
+	c.cmd.Run = c.testRunFunc()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		found := false
+		for a := range c.argsCh {
+			if a == globalFlag {
+				found = true
+				if !exp {
+					c.Errorf("got global flag as argument but expected not to: %s [%s]", c.args, globalFlag)
+				}
+			}
+		}
+
+		if !found && exp {
+			c.Errorf("did not receive flag in arguments but was expecting to: %s [%s]", c.args, globalFlag)
+		}
+	}()
+
+	Execute(c.args[1:])
+
+	wg.Wait()
+}
+
+func Test_KubectlParsing(t *testing.T) {
+	c := newCmdTest(t, clusterKubectlCmd)
+	for _, a := range [][]string{
+		{"tarmak", "kubectl", globalFlag, "arg", "arg"},
+		{"tarmak", "kubectl", "arg", globalFlag, "arg"},
+		{"tarmak", "cluster", "kubectl", globalFlag, "arg", "arg"},
+		{"tarmak", "cluster", "kubectl", "arg", globalFlag, "arg"},
+	} {
+		c.args = a
+		c.flagIgnored(true)
+	}
+
+	for _, a := range [][]string{
+		{"tarmak", globalFlag, "kubectl", "arg", "arg"},
+		{"tarmak", globalFlag, "cluster", "kubectl", "arg", "arg"},
+		{"tarmak", "cluster", globalFlag, "kubectl", "arg", "arg"},
+	} {
+		c.args = a
+		c.flagIgnored(false)
+	}
+}
+
+func Test_SSHParsing(t *testing.T) {
+	c := newCmdTest(t, clusterSshCmd)
+	for _, a := range [][]string{
+		{"tarmak", "cluster", "ssh", "arg", globalFlag, "arg"},
+		{"tarmak", "cluster", "ssh", globalFlag, "arg", "arg"},
+	} {
+		c.args = a
+		c.flagIgnored(true)
+	}
+
+	for _, a := range [][]string{
+		{"tarmak", globalFlag, "cluster", "ssh", "arg", "arg"},
+		{"tarmak", "cluster", globalFlag, "ssh", "arg", "arg"},
+	} {
+		c.args = a
+		c.flagIgnored(false)
+	}
+}

pkg/tarmak/utils/slices.go

@@ -36,3 +36,13 @@ func SliceContains(slice []string, str string) bool {
 
 	return false
 }
+
+func IndexOfString(slice []string, str string) int {
+	for i, s := range slice {
+		if s == str {
+			return i
+		}
+	}
+
+	return -1
+}