add test

cmd/keymaster/main_test.go

@@ -1,13 +1,16 @@
 package main
 
 import (
+	"bytes"
 	"crypto/ed25519"
 	"crypto/rand"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/base64"
 	"encoding/json"
+	"flag"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"os"
@@ -277,11 +280,50 @@ func TestInsertSSHCertIntoAgentORWriteToFilesystem(t *testing.T) {
 }
 
 func TestMainPrintVersion(t *testing.T) {
+	// Backup and defer restore of original os.Stdout and os.Args
+	Version = "1.0.0-test"
+	origStdout := os.Stdout
+	origArgs := os.Args
+	defer func() {
+		os.Stdout = origStdout
+		os.Args = origArgs
+	}()
+
+	// Create a pipe to capture the output
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("failed to create pipe: %v", err)
+	}
+	os.Stdout = w
+
+	// Set up the command-line arguments to include the -version flag
 	os.Args = []string{"cmd", "-version"}
+
+	// Reset the command-line flags for testing
+	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
+	printVersion = flag.Bool("version", false, "Print version and exit")
+
+	// Call the main function in a separate goroutine
 	done := make(chan struct{})
 	go func() {
 		main()
+		w.Close()
 		close(done)
 	}()
+
+	// Read the captured output
+	var buf bytes.Buffer
+	_, err = io.Copy(&buf, r)
+	if err != nil {
+		t.Fatalf("failed to read captured output: %v", err)
+	}
+
+	// Wait for the goroutine to finish
 	<-done
+
+	// Check the output
+	expectedOutput := "1.0.0-test\n"
+	if buf.String() != expectedOutput {
+		t.Errorf("expected %q, got %q", expectedOutput, buf.String())
+	}
 }