From 841d98dff064581210a370a8f93e5d80185ed144 Mon Sep 17 00:00:00 2001
From: jackgopack4 <john.peterson@datadoghq.com>
Date: Thu, 19 Sep 2024 12:24:50 -0400
Subject: [PATCH] enable test framework for versionCommand and add unit test

---
 cmd/builder/internal/command.go      |  2 +-
 cmd/builder/internal/version.go      |  5 +-
 cmd/builder/internal/version_test.go | 70 +++++++++++++++++++++++++++-
 3 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/cmd/builder/internal/command.go b/cmd/builder/internal/command.go
index 3af771787279..8015d80405c3 100644
--- a/cmd/builder/internal/command.go
+++ b/cmd/builder/internal/command.go
@@ -116,7 +116,7 @@ configuration is provided, ocb will generate a default Collector.
 		return nil, err
 	}
 	// version of this binary
-	cmd.AddCommand(versionCommand())
+	cmd.AddCommand(versionCommand(binVersion))
 
 	return cmd, nil
 }
diff --git a/cmd/builder/internal/version.go b/cmd/builder/internal/version.go
index 360a45fb4266..886a0483d8fa 100644
--- a/cmd/builder/internal/version.go
+++ b/cmd/builder/internal/version.go
@@ -15,6 +15,7 @@ var (
 )
 
 type debugReadBuildInfoFunc func() (info *debug.BuildInfo, ok bool)
+type binVersionFunc func(fn debugReadBuildInfoFunc) (string, error)
 
 // binVersion returns the version of the binary.
 // If the version is not set, it attempts to read the build information.
@@ -30,14 +31,14 @@ func binVersion(fn debugReadBuildInfoFunc) (string, error) {
 	return info.Main.Version, nil
 }
 
-func versionCommand() *cobra.Command {
+func versionCommand(fn binVersionFunc) *cobra.Command {
 	var err error
 	return &cobra.Command{
 		Use:   "version",
 		Short: "Version of ocb",
 		Long:  "Prints the version of the ocb binary",
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			version, err = binVersion(debug.ReadBuildInfo)
+			version, err = fn(debug.ReadBuildInfo)
 			if err != nil {
 				return err
 			}
diff --git a/cmd/builder/internal/version_test.go b/cmd/builder/internal/version_test.go
index 9dd66ce6d5bb..9a51d23f837a 100644
--- a/cmd/builder/internal/version_test.go
+++ b/cmd/builder/internal/version_test.go
@@ -4,9 +4,14 @@
 package internal
 
 import (
+	"bytes"
+	"context"
 	"fmt"
+	"os"
 	"runtime/debug"
 	"testing"
+
+	"github.com/spf13/cobra"
 )
 
 // Mock debug.ReadBuildInfo function
@@ -15,7 +20,7 @@ var readBuildInfo = debug.ReadBuildInfo
 func TestBinVersion(t *testing.T) {
 	// Test case: version is set
 	version = "v1.0.0"
-	v, err := binVersion(debug.ReadBuildInfo)
+	v, err := binVersion(readBuildInfo)
 	if err != nil {
 		t.Fatalf("expected no error, got %v", err)
 	}
@@ -53,3 +58,66 @@ func TestBinVersion(t *testing.T) {
 		t.Fatalf("expected empty version, got %v", v)
 	}
 }
+
+var validBinVersionFunc binVersionFunc = func(fn debugReadBuildInfoFunc) (string, error) {
+	return "v1.0.0", nil
+}
+
+var invalidBinVersionFunc binVersionFunc = func(fn debugReadBuildInfoFunc) (string, error) {
+	return "", fmt.Errorf("failed to get version")
+}
+
+func TestVersionCommand(t *testing.T) {
+	tests := []struct {
+		name           string
+		binVersion     binVersionFunc
+		expectedOutput string
+		expectedError  bool
+	}{
+		{
+			name:           "valid version",
+			binVersion:     validBinVersionFunc,
+			expectedOutput: "ocb version v1.0.0\n",
+			expectedError:  false,
+		},
+		{
+			name:           "error in binVersion",
+			binVersion:     invalidBinVersionFunc,
+			expectedOutput: "",
+			expectedError:  true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Set a mock parent command name
+			parentCmd := &cobra.Command{
+				Use: "ocb <command>",
+			}
+			// Create the command
+			var cmd = versionCommand(tt.binVersion)
+			parentCmd.AddCommand(cmd)
+			// Capture the output
+			output := bytes.NewBufferString("")
+			err_output := bytes.NewBufferString("")
+			cmd.SetOut(output)
+			cmd.SetErr(err_output)
+			// Create a new context with a fake value
+			type contextKey string
+			ctx := context.WithValue(context.Background(), contextKey("key"), "value")
+			// Set fake CLI arguments
+			fakeArgs := []string{"cmd", "version"}
+			os.Args = fakeArgs
+			// Execute the command
+			err := cmd.ExecuteContext(ctx)
+			// Check for expected error
+			if (err != nil) != tt.expectedError {
+				t.Fatalf("expected error: %v, got: %v", tt.expectedError, err)
+			}
+			// Check for expected output
+			if output.String() != tt.expectedOutput {
+				t.Fatalf("expected output: %v, got: %v", tt.expectedOutput, output.String())
+			}
+		})
+	}
+}