test: improve CLI testing

Signed-off-by: Andres Taylor <[email protected]>

go/cmd/keys.go

@@ -35,7 +35,7 @@ func keysCmd() *cobra.Command {
 		PreRun: func(cmd *cobra.Command, _ []string) {
 			csvConfig = csvFlagsToConfig(cmd, *flags)
 		},
-		RunE: func(_ *cobra.Command, args []string) error {
+		RunE: func(c *cobra.Command, args []string) error {
 			cfg := keys.Config{
 				FileName: args[0],
 			}
@@ -46,7 +46,7 @@ func keysCmd() *cobra.Command {
 			}
 			cfg.Loader = loader
 
-			return keys.Run(cfg)
+			return keys.Run(c.OutOrStdout(), cfg)
 		},
 	}
 

go/cmd/main_test.go

@@ -0,0 +1,57 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package cmd
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestApp(t *testing.T) {
+	tests := []struct {
+		args []string
+		want string
+	}{
+		{[]string{"help"}, "Utils tools for testing, running and benchmarking Vitess"},
+		{[]string{"keys", "../../t/demo.test"}, `"queryStructure"`},
+		{[]string{"keys", "--help"}, `Runs vexplain keys on all queries of the test file`},
+		{[]string{"txs", "--help"}, `Analyze transactions on a query log`},
+		{[]string{"test", "--help"}, `Test the given workload against both Vitess and MySQL`},
+		{[]string{"trace", "--help"}, `Runs the given workload and does a`},
+		{[]string{"summarize", "--help"}, `Compares and analyses a trace output`},
+		{[]string{"dbinfo", "--help"}, `Loads info from the database including row counts`},
+		{[]string{"planalyze", "--help"}, `Analyze the query plan`},
+	}
+	for _, tt := range tests {
+		t.Run("vt  "+strings.Join(tt.args, " "), func(t *testing.T) {
+			cmd := getRootCmd()
+			buf := new(bytes.Buffer)
+			cmd.SetOut(buf)
+			cmd.SetErr(buf)
+			cmd.SetArgs(tt.args)
+
+			err := cmd.Execute()
+
+			out := buf.String()
+			require.NoError(t, err, out)
+			require.Contains(t, out, tt.want)
+		})
+	}
+}

go/cmd/root.go

@@ -25,6 +25,13 @@ import (
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
+	err := getRootCmd().Execute()
+	if err != nil {
+		os.Exit(1)
+	}
+}
+
+func getRootCmd() *cobra.Command {
 	// rootCmd represents the base command when called without any subcommands
 	root := &cobra.Command{
 		Use:   "vt",
@@ -39,8 +46,5 @@ func Execute() {
 	root.AddCommand(dbinfoCmd())
 	root.AddCommand(transactionsCmd())
 	root.AddCommand(planalyzeCmd())
-	err := root.Execute()
-	if err != nil {
-		os.Exit(1)
-	}
+	return root
 }

go/keys/keys.go

@@ -21,7 +21,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"os"
 	"sort"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -73,11 +72,7 @@ type (
 	}
 )
 
-func Run(cfg Config) error {
-	return run(os.Stdout, cfg)
-}
-
-func run(out io.Writer, cfg Config) error {
+func Run(out io.Writer, cfg Config) error {
 	si := &SchemaInfo{
 		Tables: make(map[string]Columns),
 	}

go/keys/keys_test.go

@@ -65,7 +65,7 @@ func TestKeys(t *testing.T) {
 	for _, tcase := range cases {
 		t.Run(tcase.expectedFile, func(t *testing.T) {
 			sb := &strings.Builder{}
-			err := run(sb, tcase.cfg)
+			err := Run(sb, tcase.cfg)
 			require.NoError(t, err)
 
 			out, err := os.ReadFile("../testdata/keys-output/" + tcase.expectedFile)