From b99cc783899589a285c972bd27a8f96fcc3a9ba6 Mon Sep 17 00:00:00 2001
From: Brian Pursley <bpursley@cinlogic.com>
Date: Thu, 12 Mar 2020 19:27:27 -0400
Subject: [PATCH] Add custom usage template for cobra command so that usage
 will display 'kubectl krew' instead of just 'krew' See issue #521 for more
 details: https://github.com/kubernetes-sigs/krew/issues/521

---
 .gitignore                |  4 ++++
 cmd/krew/cmd/root.go      |  9 +++++++-
 cmd/krew/cmd/root_test.go | 43 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 cmd/krew/cmd/root_test.go

diff --git a/.gitignore b/.gitignore
index 2de4e72f..b1e608c4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,6 +11,10 @@
 # Output of the go coverage tool, specifically when used with LiteIDE
 *.out
 
+# Files generated by JetBrains IDEs, e.g. IntelliJ IDEA
+.idea/
+*.iml
+
 out/
 build/
 coverage.txt
diff --git a/cmd/krew/cmd/root.go b/cmd/krew/cmd/root.go
index 82c6104c..0fc84213 100644
--- a/cmd/krew/cmd/root.go
+++ b/cmd/krew/cmd/root.go
@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"math/rand"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/fatih/color"
@@ -58,7 +59,7 @@ var (
 
 // rootCmd represents the base command when called without any subcommands
 var rootCmd = &cobra.Command{
-	Use:   "krew",
+	Use:   "krew", // This is prefixed by kubectl in the custom usage template
 	Short: "krew is the kubectl plugin manager",
 	Long: `krew is the kubectl plugin manager.
 You can invoke krew through kubectl: "kubectl krew [command]..."`,
@@ -98,6 +99,12 @@ func init() {
 	}
 
 	paths = environment.MustGetKrewPaths()
+
+	// Cobra doesn't have a way to specify a two word command (ie. "kubectl krew"), so set a custom usage template
+	// with kubectl in it. Cobra will use this template for the root and all child commands.
+	rootCmd.SetUsageTemplate(strings.NewReplacer(
+		"{{.UseLine}}", "kubectl {{.UseLine}}",
+		"{{.CommandPath}}", "kubectl {{.CommandPath}}").Replace(rootCmd.UsageTemplate()))
 }
 
 func preRun(cmd *cobra.Command, _ []string) error {
diff --git a/cmd/krew/cmd/root_test.go b/cmd/krew/cmd/root_test.go
new file mode 100644
index 00000000..8ab6c0b0
--- /dev/null
+++ b/cmd/krew/cmd/root_test.go
@@ -0,0 +1,43 @@
+// Copyright 2020 The Kubernetes 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 (
+	"strings"
+	"testing"
+)
+
+func TestUsageTemplateContainsReplacements(t *testing.T) {
+	usageTemplate := rootCmd.UsageTemplate()
+
+	expectedStrings := []string{
+		"kubectl {{.CommandPath}}",
+		"kubectl {{.UseLine}}",
+	}
+
+	for _, expectedString := range expectedStrings {
+		if !strings.Contains(usageTemplate, expectedString) {
+			t.Errorf("Expected usage template to contain '%s' but it did not:\n%v", expectedString, usageTemplate)
+		}
+	}
+}
+
+func TestUsageStringContainsKubectlKrew(t *testing.T) {
+	usage := rootCmd.UsageString()
+
+	if !strings.Contains(usage, "kubectl krew") {
+		t.Errorf("Expected usage string to contain 'kubectl krew' but it did not:\n%v", usage)
+	}
+}