test cobra

cmd/api/main.go

@@ -8,36 +8,24 @@ import (
 	"OpenIM/pkg/common/mw"
 	"fmt"
 	"github.com/OpenIMSDK/openKeeper"
-	"github.com/spf13/cobra"
 	"os"
 	"strconv"
 
 	"OpenIM/pkg/common/constant"
 )
 
-var startCmd = &cobra.Command{
-	Use:   "start",
-	Short: "Start the server",
-	Run: func(cmd *cobra.Command, args []string) {
-		port, _ := cmd.Flags().GetInt(constant.FlagPort)
-		if err := run(port); err != nil {
-			panic(err.Error())
-		}
-	},
-}
-
 func main() {
-	startCmd.Flags().IntP(constant.FlagPort, "p", 0, "Port to listen on")
-	startCmd.Flags().StringP(constant.FlagConf, "c", "", "Path to config file folder")
-	rootCmd := cmd.NewRootCmd(nil)
-	rootCmd.Command.AddCommand(startCmd)
-	if err := startCmd.Execute(); err != nil {
+	rootCmd := cmd.NewRootCmd()
+	rootCmd.AddPortFlag()
+	rootCmd.AddRunE(run)
+	if err := rootCmd.Command.Execute(); err != nil {
 		fmt.Println(err)
 		os.Exit(1)
 	}
 }
 
-func run(port int) error {
+func run(rootCmd cmd.RootCmd) error {
+	port := rootCmd.GetPortFlag()
 	if port == 0 {
 		port = config.Config.Api.GinPort[0]
 	}

pkg/common/cmd/root.go

@@ -7,29 +7,69 @@ import (
 )
 
 type RootCmd struct {
-	Command cobra.Command
+	Command  cobra.Command
+	port     int
+	portFlag bool
+
+	prometheusPort     int
+	prometheusPortFlag bool
 }
 
-func NewRootCmd(f func() error) RootCmd {
+func NewRootCmd() RootCmd {
 	c := cobra.Command{
 		Use:   "start",
 		Short: "Start the server",
 		Long:  `Start the server`,
-		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
-			return getConfFromCmdAndInit(cmd)
-		},
 	}
-	c.Flags()
-	rootCmd := RootCmd{Command: c}
+	rootCmd := RootCmd{}
+	c.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
+		if rootCmd.portFlag {
+			rootCmd.port = rootCmd.getPortFlag(cmd)
+		}
+		if rootCmd.prometheusPortFlag {
+			rootCmd.prometheusPort = rootCmd.GetPrometheusPortFlag(cmd)
+		}
+		return rootCmd.getConfFromCmdAndInit(cmd)
+	}
 	rootCmd.init()
 	return rootCmd
 }
 
+func (r RootCmd) AddRunE(f func(cmd RootCmd) error) {
+	r.Command.RunE = func(cmd *cobra.Command, args []string) error {
+		return f(r)
+	}
+}
+
 func (r RootCmd) init() {
 	r.Command.Flags().StringP(constant.FlagConf, "c", "", "Path to config file folder")
 }
 
-func getConfFromCmdAndInit(cmdLines *cobra.Command) error {
+func (r RootCmd) AddPortFlag() {
+	r.Command.Flags().StringP(constant.FlagPort, "p", "", "server listen port")
+	r.portFlag = true
+}
+
+func (r RootCmd) getPortFlag(cmd *cobra.Command) int {
+	port, _ := cmd.Flags().GetInt(constant.FlagPort)
+	return port
+}
+
+func (r RootCmd) GetPortFlag() int {
+	return r.port
+}
+
+func (r RootCmd) AddPrometheusPortFlag() {
+	r.Command.Flags().StringP(constant.PrometheusPort, "pp", "", "server listen port")
+	r.prometheusPortFlag = true
+}
+
+func (r RootCmd) GetPrometheusPortFlag(cmd *cobra.Command) int {
+	port, _ := cmd.Flags().GetInt(constant.PrometheusPort)
+	return port
+}
+
+func (r RootCmd) getConfFromCmdAndInit(cmdLines *cobra.Command) error {
 	configFolderPath, _ := cmdLines.Flags().GetString(constant.FlagConf)
 	return config.InitConfig(configFolderPath)
 }