From 4cb9eec25726f3f2502713cb576e6c50bb4277de Mon Sep 17 00:00:00 2001
From: Brendan Le Glaunec <brendan.le-glaunec@epitech.eu>
Date: Mon, 2 Dec 2019 17:34:06 +0100
Subject: [PATCH] Add custom help function to command

---
 pkg/cli/commands.go      | 29 +++++++++++++------
 pkg/cli/commands_test.go | 61 ++++++++++++++++++++++++++++++++++++++++
 pkg/cli/loader_file.go   |  2 +-
 3 files changed, 82 insertions(+), 10 deletions(-)

diff --git a/pkg/cli/commands.go b/pkg/cli/commands.go
index 00ad09c8b..8151f3c26 100644
--- a/pkg/cli/commands.go
+++ b/pkg/cli/commands.go
@@ -4,18 +4,20 @@ package cli
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 )
 
 // Command structure contains program/command information (command name and description).
 type Command struct {
-	Name          string
-	Description   string
-	Configuration interface{}
-	Resources     []ResourceLoader
-	Run           func([]string) error
-	Hidden        bool
+	Name           string
+	Description    string
+	Configuration  interface{}
+	Resources      []ResourceLoader
+	Run            func([]string) error
+	CustomHelpFunc func(io.Writer, *Command) error
+	Hidden         bool
 	// AllowArg if not set, disallows any argument that is not a known command or a sub-command.
 	AllowArg    bool
 	subCommands []*Command
@@ -35,6 +37,15 @@ func (c *Command) AddCommand(cmd *Command) error {
 	return nil
 }
 
+// PrintHelp calls the custom help function of the command if it's set.
+// Otherwise, it calls the default help function.
+func (c *Command) PrintHelp(w io.Writer) error {
+	if c.CustomHelpFunc != nil {
+		return c.CustomHelpFunc(w, c)
+	}
+	return PrintHelp(w, c)
+}
+
 // Execute Executes a command.
 func Execute(cmd *Command) error {
 	return execute(cmd, os.Args, true)
@@ -92,16 +103,16 @@ func execute(cmd *Command, args []string, root bool) error {
 
 func run(cmd *Command, args []string) error {
 	if len(args) > 0 && !isFlag(args[0]) && !cmd.AllowArg {
-		_ = PrintHelp(os.Stdout, cmd)
+		_ = cmd.PrintHelp(os.Stdout)
 		return fmt.Errorf("command not found: %s", args[0])
 	}
 
 	if isHelp(args) {
-		return PrintHelp(os.Stdout, cmd)
+		return cmd.PrintHelp(os.Stdout)
 	}
 
 	if cmd.Run == nil {
-		_ = PrintHelp(os.Stdout, cmd)
+		_ = cmd.PrintHelp(os.Stdout)
 		return fmt.Errorf("command %s is not runnable", cmd.Name)
 	}
 
diff --git a/pkg/cli/commands_test.go b/pkg/cli/commands_test.go
index 4a1b9a0ac..39fa80a34 100644
--- a/pkg/cli/commands_test.go
+++ b/pkg/cli/commands_test.go
@@ -1,6 +1,10 @@
 package cli
 
 import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"strings"
@@ -55,6 +59,63 @@ func TestCommand_AddCommand(t *testing.T) {
 	}
 }
 
+func TestCommand_PrintHelp(t *testing.T) {
+	testCases := []struct {
+		desc           string
+		command        *Command
+		expectedOutput string
+		expectedError  error
+	}{
+		{
+			desc:           "print default help",
+			command:        &Command{},
+			expectedOutput: "    \n\nUsage:  [command] [flags] [arguments]\n\nUse \" [command] --help\" for help on any command.\n\n",
+		},
+		{
+			desc: "print custom help",
+			command: &Command{
+				Name:        "root",
+				Description: "Description for root",
+				Configuration: &struct {
+					Foo []struct {
+						Field string
+					}
+				}{},
+				Run: func(args []string) error {
+					return nil
+				},
+				CustomHelpFunc: func(w io.Writer, _ *Command) error {
+					_, _ = fmt.Fprintln(w, "test")
+					return nil
+				},
+			},
+			expectedOutput: "test\n",
+		},
+		{
+			desc: "error is returned from called help",
+			command: &Command{
+				CustomHelpFunc: func(_ io.Writer, _ *Command) error {
+					return errors.New("test")
+				},
+			},
+			expectedError: errors.New("test"),
+		},
+	}
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			buffer := &bytes.Buffer{}
+			err := test.command.PrintHelp(buffer)
+
+			assert.Equal(t, test.expectedError, err)
+			assert.Equal(t, test.expectedOutput, buffer.String())
+		})
+	}
+}
+
 func Test_execute(t *testing.T) {
 	var called string
 
diff --git a/pkg/cli/loader_file.go b/pkg/cli/loader_file.go
index eb0b3f78e..d6ad58b7b 100644
--- a/pkg/cli/loader_file.go
+++ b/pkg/cli/loader_file.go
@@ -25,7 +25,7 @@ func (f *FileLoader) GetFilename() string {
 func (f *FileLoader) Load(args []string, cmd *Command) (bool, error) {
 	ref, err := flag.Parse(args, cmd.Configuration)
 	if err != nil {
-		_ = PrintHelp(os.Stdout, cmd)
+		_ = cmd.PrintHelp(os.Stdout)
 		return false, err
 	}