From 23a6174c7f9d0c64a48e634054bdf6886fb7fbba Mon Sep 17 00:00:00 2001
From: Joon-Ho Son <20669173+sonjoonho@users.noreply.github.com>
Date: Thu, 21 Jan 2021 03:33:16 +0000
Subject: [PATCH] Add the ability to specify a filePostpender in
 GenMarkdownTreeCustom (#1270)

---
 doc/md_docs.go      | 11 +++++++----
 doc/md_docs_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/doc/md_docs.go b/doc/md_docs.go
index 5181af8..fafab26 100644
--- a/doc/md_docs.go
+++ b/doc/md_docs.go
@@ -122,17 +122,17 @@ func GenMarkdownCustom(cmd *cobra.Command, w io.Writer, linkHandler func(string)
 func GenMarkdownTree(cmd *cobra.Command, dir string) error {
 	identity := func(s string) string { return s }
 	emptyStr := func(s string) string { return "" }
-	return GenMarkdownTreeCustom(cmd, dir, emptyStr, identity)
+	return GenMarkdownTreeCustom(cmd, dir, emptyStr, emptyStr, identity)
 }
 
 // GenMarkdownTreeCustom is the the same as GenMarkdownTree, but
-// with custom filePrepender and linkHandler.
-func GenMarkdownTreeCustom(cmd *cobra.Command, dir string, filePrepender, linkHandler func(string) string) error {
+// with custom filePrepender, filePostpender, and linkHandler.
+func GenMarkdownTreeCustom(cmd *cobra.Command, dir string, filePrepender, filePostpender, linkHandler func(string) string) error {
 	for _, c := range cmd.Commands() {
 		if !c.IsAvailableCommand() || c.IsAdditionalHelpTopicCommand() {
 			continue
 		}
-		if err := GenMarkdownTreeCustom(c, dir, filePrepender, linkHandler); err != nil {
+		if err := GenMarkdownTreeCustom(c, dir, filePrepender, filePostpender, linkHandler); err != nil {
 			return err
 		}
 	}
@@ -151,5 +151,8 @@ func GenMarkdownTreeCustom(cmd *cobra.Command, dir string, filePrepender, linkHa
 	if err := GenMarkdownCustom(cmd, f, linkHandler); err != nil {
 		return err
 	}
+	if _, err := io.WriteString(f, filePostpender(filename)); err != nil {
+		return err
+	}
 	return nil
 }
diff --git a/doc/md_docs_test.go b/doc/md_docs_test.go
index f325167..29430b4 100644
--- a/doc/md_docs_test.go
+++ b/doc/md_docs_test.go
@@ -95,6 +95,37 @@ func TestGenMdTree(t *testing.T) {
 	}
 }
 
+func TestGenMdTreeCustom(t *testing.T) {
+	tmpdir, err := ioutil.TempDir("", "test-gen-md-tree")
+	if err != nil {
+		t.Fatalf("Failed to create tmpdir: %v", err)
+	}
+	defer os.RemoveAll(tmpdir)
+
+	prepender := func(s string) string { return "Prepended" }
+	postpender := func(s string) string { return "Postpended" }
+	identity := func(s string) string { return s }
+
+	if err := GenMarkdownTreeCustom(rootCmd, tmpdir, prepender, postpender, identity); err != nil {
+		t.Fatalf("GenMarkdownTree failed: %v", err)
+	}
+
+	gotRoot := fileContents(t, tmpdir, "root.md")
+	checkStringContains(t, gotRoot, "Prepended")
+	checkStringContains(t, gotRoot, rootCmd.Long)
+	checkStringContains(t, gotRoot, "Postpended")
+
+	gotEcho := fileContents(t, tmpdir, "root_echo.md")
+	checkStringContains(t, gotEcho, "Prepended")
+	checkStringContains(t, gotEcho, echoCmd.Long)
+	checkStringContains(t, gotEcho, "Postpended")
+
+	gotEchoSub := fileContents(t, tmpdir, "root_echo_echosub.md")
+	checkStringContains(t, gotEchoSub, "Prepended")
+	checkStringContains(t, gotEchoSub, echoSubCmd.Long)
+	checkStringContains(t, gotEchoSub, "Postpended")
+}
+
 func BenchmarkGenMarkdownToFile(b *testing.B) {
 	file, err := ioutil.TempFile("", "")
 	if err != nil {
@@ -110,3 +141,11 @@ func BenchmarkGenMarkdownToFile(b *testing.B) {
 		}
 	}
 }
+
+func fileContents(t *testing.T, dir, filename string) string {
+	contents, err := ioutil.ReadFile(filepath.Join(dir, filename))
+	if err != nil {
+		t.Fatalf("Error loading file %q: %v ", filename, err)
+	}
+	return string(contents)
+}