Add tests for Execute

cmd/root/root.go

@@ -92,9 +92,10 @@ func flagErrorFunc(c *cobra.Command, err error) error {
 
 // 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(cmd *cobra.Command) {
+//
+// Returns exit code for os.Exit.
+func Execute(ctx context.Context, cmd *cobra.Command) int {
 	// TODO: deferred panic recovery
-	ctx := context.Background()
 
 	// Run the command
 	cmd, err := cmd.ExecuteContextC(ctx)
@@ -119,6 +120,8 @@ func Execute(cmd *cobra.Command) {
 	}
 
 	if err != nil {
-		os.Exit(1)
+		return 1
+	} else {
+		return 0
 	}
 }

main.go

@@ -2,11 +2,15 @@ package main
 
 import (
 	"context"
+	"os"
 
 	"github.com/databricks/cli/cmd"
 	"github.com/databricks/cli/cmd/root"
 )
 
 func main() {
-	root.Execute(cmd.New(context.Background()))
+	ctx := context.Background()
+	code := root.Execute(ctx, cmd.New(ctx))
+
+	os.Exit(code)
 }

main_test.go

@@ -1,7 +1,12 @@
 package main
 
 import (
+	"bufio"
+	"bytes"
 	"context"
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/flags"
 	"testing"
 
 	"github.com/databricks/cli/cmd"
@@ -23,3 +28,45 @@ func TestCommandsDontUseUnderscoreInName(t *testing.T) {
 		queue = append(queue[1:], cmd.Commands()...)
 	}
 }
+
+func TestExecute_version(t *testing.T) {
+	stderr, logger := createFakeLogger()
+	stdout := bytes.Buffer{}
+	ctx := cmdio.NewContext(context.Background(), logger)
+
+	cli := cmd.New(ctx)
+	cli.SetArgs([]string{"version"})
+	cli.SetOut(&stdout)
+
+	code := root.Execute(ctx, cli)
+
+	assert.Equal(t, "", string(stderr.Bytes()))
+	assert.Contains(t, stdout.String(), "Databricks CLI v")
+	assert.Equal(t, 0, code)
+}
+
+func TestExecute_unknownCommand(t *testing.T) {
+	stderr, logger := createFakeLogger()
+	stdout := bytes.Buffer{}
+	ctx := cmdio.NewContext(context.Background(), logger)
+
+	cli := cmd.New(ctx)
+	cli.SetOut(&stdout)
+	cli.SetArgs([]string{"abcabcabc"})
+
+	code := root.Execute(ctx, cli)
+
+	assert.Equal(t, `Error: unknown command "abcabcabc" for "databricks"`+"\n", stderr.String())
+	assert.Equal(t, "", stdout.String())
+	assert.Equal(t, 1, code)
+}
+
+func createFakeLogger() (*bytes.Buffer, *cmdio.Logger) {
+	out := bytes.Buffer{}
+
+	return &out, &cmdio.Logger{
+		Mode:   flags.ModeAppend,
+		Reader: bufio.Reader{},
+		Writer: &out,
+	}
+}