chore: add version flag test

Author: Fleezesd

pkg/version/flag.go

@@ -10,10 +10,13 @@ package version
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"strconv"
+	"strings"
 
 	flag "github.com/spf13/pflag"
+	"k8s.io/component-base/version"
 )
 
 type versionValue int
@@ -39,6 +42,16 @@ func (v *versionValue) Set(s string) error {
 		*v = VersionRaw
 		return nil
 	}
+
+	if strings.HasPrefix(s, "v") {
+		err := version.SetDynamicVersion(s)
+		if err == nil {
+			value, _ := strconv.Atoi(s)
+			*v = versionValue(value)
+		}
+		return err
+	}
+
 	boolVal, err := strconv.ParseBool(s)
 	if boolVal {
 		*v = VersionTrue
@@ -83,14 +96,21 @@ func AddFlags(fs *flag.FlagSet) {
 	fs.AddFlag(flag.Lookup(versionFlagName))
 }
 
+// variables for unit testing PrintAndExitIfRequested
+var (
+	appName = "onex-apiserver"
+	output  = io.Writer(os.Stdout)
+	exit    = os.Exit
+)
+
 // PrintAndExitIfRequested will check if the -version flag was passed
 // and, if so, print the version and exit.
 func PrintAndExitIfRequested(appName string) {
 	if *versionFlag == VersionRaw {
-		fmt.Printf("%s\n", Get().Text())
-		os.Exit(0)
+		fmt.Fprintf(output, "%s\n", Get().Text())
+		exit(0)
 	} else if *versionFlag == VersionTrue {
-		fmt.Printf("%s %s\n", appName, Get().GitVersion)
-		os.Exit(0)
+		fmt.Fprintf(output, "%s %s\n", appName, Get().GitVersion)
+		exit(0)
 	}
 }

pkg/version/flag_test.go

@@ -0,0 +1,140 @@
+package version
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/spf13/pflag"
+	"k8s.io/component-base/version"
+)
+
+func TestVersionFlag(t *testing.T) {
+	initialFlagValue := int(*versionFlag)
+	initialVersion := Get()
+
+	testcases := []struct {
+		name               string
+		flags              []string
+		expectError        string
+		expectExit         bool
+		expectPrintVersion string
+		expectGitVersion   string
+	}{
+		{
+			name:             "no flag",
+			flags:            []string{},
+			expectGitVersion: initialVersion.GitVersion,
+		},
+		{
+			name:             "false",
+			flags:            []string{"--version=false"},
+			expectGitVersion: initialVersion.GitVersion,
+		},
+
+		{
+			name:               "valueless",
+			flags:              []string{"--version"},
+			expectGitVersion:   initialVersion.GitVersion,
+			expectExit:         true,
+			expectPrintVersion: appName + " " + initialVersion.GitVersion,
+		},
+		{
+			name:               "true",
+			flags:              []string{"--version=true"},
+			expectGitVersion:   initialVersion.GitVersion,
+			expectExit:         true,
+			expectPrintVersion: appName + " " + initialVersion.GitVersion,
+		},
+		{
+			name:               "raw",
+			flags:              []string{"--version=raw"},
+			expectGitVersion:   initialVersion.GitVersion,
+			expectExit:         true,
+			expectPrintVersion: fmt.Sprintf("%s", strings.TrimSpace(initialVersion.Text())),
+		},
+		{
+			name:               "truthy",
+			flags:              []string{"--version=T"},
+			expectGitVersion:   initialVersion.GitVersion,
+			expectExit:         true,
+			expectPrintVersion: appName + " " + initialVersion.GitVersion,
+		},
+		{
+			name:             "override",
+			flags:            []string{"--version=v0.0.0-custom"},
+			expectGitVersion: "v0.0.0-custom",
+		},
+		{
+			name:        "invalid override semver",
+			flags:       []string{"--version=vX"},
+			expectError: `could not parse "vX"`,
+		},
+		{
+			name:        "invalid override major",
+			flags:       []string{"--version=v1.0.0"},
+			expectError: `must match major/minor/patch`,
+		},
+		{
+			name:        "invalid override minor",
+			flags:       []string{"--version=v0.1.0"},
+			expectError: `must match major/minor/patch`,
+		},
+		{
+			name:        "invalid override patch",
+			flags:       []string{"--version=v0.0.1"},
+			expectError: `must match major/minor/patch`,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+
+			originalOutput := output
+			originalExit := exit
+
+			outputBuffer := &bytes.Buffer{}
+			output = outputBuffer
+			exitCalled := false
+			exit = func(code int) { exitCalled = true }
+
+			t.Cleanup(func() {
+				output = originalOutput
+				exit = originalExit
+				*versionFlag = versionValue(initialFlagValue)
+				err := version.SetDynamicVersion(initialVersion.GitVersion)
+				if err != nil {
+					t.Fatal(err)
+				}
+			})
+
+			fs := pflag.NewFlagSet("test", pflag.ContinueOnError)
+			AddFlags(fs)
+			err := fs.Parse(tc.flags)
+			if tc.expectError != "" {
+				if err == nil {
+					t.Fatal("expected error, got none")
+				}
+				if !strings.Contains(err.Error(), tc.expectError) {
+					t.Fatalf("expected error containing %q, got %q", tc.expectError, err.Error())
+				}
+				return
+			} else if err != nil {
+				t.Fatalf("unexpected parse error: %v", err)
+			}
+
+			if e, a := tc.expectGitVersion, version.Get().GitVersion; e != a {
+				t.Fatalf("gitversion: expected %v, got %v", e, a)
+			}
+
+			PrintAndExitIfRequested(appName)
+			if e, a := tc.expectExit, exitCalled; e != a {
+				t.Fatalf("exit(): expected %v, got %v", e, a)
+			}
+			if e, a := tc.expectPrintVersion, strings.TrimSpace(outputBuffer.String()); e != a {
+				t.Fatalf("print version: expected %v, got %v", e, a)
+			}
+		})
+	}
+}