From 525cf0b9d58bf5e8b9cf4474b4dcba1fb44496f7 Mon Sep 17 00:00:00 2001
From: Marc Khouzam <marc.khouzam@montreal.ca>
Date: Fri, 30 Oct 2020 09:53:10 -0400
Subject: [PATCH] feat(comp): Completion for the docs --type flag

Add completion for the --type flag of the docs command.

Signed-off-by: Marc Khouzam <marc.khouzam@montreal.ca>
---
 cmd/helm/docs.go                                    | 11 +++++++++++
 cmd/helm/docs_test.go                               | 13 +++++++++++++
 cmd/helm/testdata/output/docs-type-comp.txt         |  5 +++++
 .../testdata/output/docs-type-filtered-comp.txt     |  3 +++
 4 files changed, 32 insertions(+)
 create mode 100644 cmd/helm/testdata/output/docs-type-comp.txt
 create mode 100644 cmd/helm/testdata/output/docs-type-filtered-comp.txt

diff --git a/cmd/helm/docs.go b/cmd/helm/docs.go
index df558d3e8..1a28a47ec 100644
--- a/cmd/helm/docs.go
+++ b/cmd/helm/docs.go
@@ -68,6 +68,17 @@ func newDocsCmd(out io.Writer) *cobra.Command {
 	f.StringVar(&o.docTypeString, "type", "markdown", "the type of documentation to generate (markdown, man, bash)")
 	f.BoolVar(&o.generateHeaders, "generate-headers", false, "generate standard headers for markdown files")
 
+	cmd.RegisterFlagCompletionFunc("type", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		types := []string{"bash", "man", "markdown"}
+		var comps []string
+		for _, t := range types {
+			if strings.HasPrefix(t, toComplete) {
+				comps = append(comps, t)
+			}
+		}
+		return comps, cobra.ShellCompDirectiveNoFileComp
+	})
+
 	return cmd
 }
 
diff --git a/cmd/helm/docs_test.go b/cmd/helm/docs_test.go
index cda299b0a..f0082578a 100644
--- a/cmd/helm/docs_test.go
+++ b/cmd/helm/docs_test.go
@@ -20,6 +20,19 @@ import (
 	"testing"
 )
 
+func TestDocsTypeFlagCompletion(t *testing.T) {
+	tests := []cmdTestCase{{
+		name:   "completion for docs --type",
+		cmd:    "__complete docs --type ''",
+		golden: "output/docs-type-comp.txt",
+	}, {
+		name:   "completion for docs --type",
+		cmd:    "__complete docs --type mar",
+		golden: "output/docs-type-filtered-comp.txt",
+	}}
+	runTestCmd(t, tests)
+}
+
 func TestDocsFileCompletion(t *testing.T) {
 	checkFileCompletion(t, "docs", false)
 }
diff --git a/cmd/helm/testdata/output/docs-type-comp.txt b/cmd/helm/testdata/output/docs-type-comp.txt
new file mode 100644
index 000000000..69494f87d
--- /dev/null
+++ b/cmd/helm/testdata/output/docs-type-comp.txt
@@ -0,0 +1,5 @@
+bash
+man
+markdown
+:4
+Completion ended with directive: ShellCompDirectiveNoFileComp
diff --git a/cmd/helm/testdata/output/docs-type-filtered-comp.txt b/cmd/helm/testdata/output/docs-type-filtered-comp.txt
new file mode 100644
index 000000000..55104f32e
--- /dev/null
+++ b/cmd/helm/testdata/output/docs-type-filtered-comp.txt
@@ -0,0 +1,3 @@
+markdown
+:4
+Completion ended with directive: ShellCompDirectiveNoFileComp