mukan-sdk/internal/testutil/cmd.go

package testutil

import (
	"strings"
	"testing"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// ResetArgs resets arguments for the command. It is desired to be a helpful function for the cmd.SetArgs
// in the case of calling multiple times in a unit test, as cmd.SetArgs doesn't
// reset the flag value as expected.
//
// **Warning**: this is only compatible with following flag types:
//  1. the implementations of pflag.Value
//  2. the built-in implementations of pflag.SliceValue
//  3. the custom implementations of pflag.SliceValue that are split by comma ","
//
// see https://github.com/spf13/cobra/issues/2079#issuecomment-1870115781 for more detail info
func ResetArgs(t *testing.T, cmd *cobra.Command) {
	t.Helper()
	// if flags haven't been parsed yet, there is no need to reset the args
	if !cmd.Flags().Parsed() {
		return
	}
	// If flags have already been parsed, we should reset the values of flags that haven't been set
	cmd.Flags().Visit(func(pf *pflag.Flag) {
		// if the flag hasn't been changed, there is no need to reset the args
		if !pf.Changed {
			return
		}
		// handle pflag.SliceValue
		if v, ok := pf.Value.(pflag.SliceValue); ok {
			defVal := strings.Trim(pf.DefValue, "[]")
			defSliceVal := make([]string, 0)
			if defVal != "" {
				defSliceVal = strings.Split(defVal, ",")
			}
			if err := v.Replace(defSliceVal); err != nil {
				t.Errorf("error resetting argument <%s> with default value <%+v>: %v", pf.Name, defSliceVal, err)
			}
			return
		}
		// handle pflag.Value
		if err := pf.Value.Set(pf.DefValue); err != nil {
			t.Errorf("error resetting argument <%s> with default value <%s>: %v", pf.Name, pf.DefValue, err)
		}
	})
}