Skip to content

Commit

Permalink
Call flag.Parse() to parse global flags like -logtostderr (#198)
Browse files Browse the repository at this point in the history
  • Loading branch information
zbud-msft authored Mar 12, 2024
1 parent 736e3b4 commit 758ec18
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 10 deletions.
54 changes: 53 additions & 1 deletion telemetry/telemetry.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"time"
"os"
"os/signal"
"strings"
"syscall"
"sync"
log "github.com/golang/glog"
Expand Down Expand Up @@ -50,8 +51,20 @@ func main() {
}

func runTelemetry(args []string) error {
/* Glog flags like -logtostderr have to be part of the global flagset.
Because we use a custom flagset to avoid the use of global var and improve
testability, we have to parse cmd line args two different times such that
in the first parse, cmd line args will contain global flags and flag.Parse() will be called.
The second parse, cmd line args will contain flags only relevant to telemetry, and our custom flagset will
call Parse().
*/
glogFlags, telemetryFlags := parseOSArgs()
os.Args = glogFlags
flag.Parse() // glog flags will be populated after global flag parse

os.Args = telemetryFlags
fs := flag.NewFlagSet(args[0], flag.ExitOnError)
telemetryCfg, cfg, err := setupFlags(fs)
telemetryCfg, cfg, err := setupFlags(fs) // telemetry flags will be populated after second parse
if err != nil {
return err
}
Expand All @@ -72,6 +85,45 @@ func runTelemetry(args []string) error {
return nil
}

func getGlogFlagsMap() map[string] bool {
// glog flags: https://pkg.go.dev/github.com/golang/glog
return map[string]bool {
"-alsologtostderr": true,
"-log_backtrace_at": true,
"-log_dir": true,
"-log_link": true,
"-logbuflevel": true,
"-logtostderr": true,
"-stderrthreshold": true,
"-v": true,
"-vmodule": true,
}
}

func parseOSArgs() ([]string, []string) {
glogFlags := []string{os.Args[0]}
telemetryFlags := []string{os.Args[0]}
glogFlagsMap := getGlogFlagsMap()

for _, arg := range os.Args[1:] {
if strings.HasPrefix(arg, "-v") { // only flag in both glog and telemetry
glogFlags = append(glogFlags, arg)
telemetryFlags = append(telemetryFlags, arg)
continue
}
isGlogFlag := false
if glogFlagsMap[arg] {
glogFlags = append(glogFlags, arg)
isGlogFlag = true
continue
}
if !isGlogFlag {
telemetryFlags = append(telemetryFlags, arg)
}
}
return glogFlags, telemetryFlags
}

func setupFlags(fs *flag.FlagSet) (*TelemetryConfig, *gnmi.Config, error) {
telemetryCfg := &TelemetryConfig {
UserAuth: gnmi.AuthTypes{"password": false, "cert": false, "jwt": false},
Expand Down
21 changes: 12 additions & 9 deletions telemetry/telemetry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,25 +23,28 @@ import (
)

func TestRunTelemetry(t *testing.T) {
telemetryCfg := &TelemetryConfig{}
cfg := &gnmi.Config{}

patches := gomonkey.ApplyFunc(setupFlags, func(*flag.FlagSet) (*TelemetryConfig, *gnmi.Config, error) {
return telemetryCfg, cfg, nil
})
patches.ApplyFunc(startGNMIServer, func(_ *TelemetryConfig, _ *gnmi.Config, serverControlSignal <-chan int, wg *sync.WaitGroup) {
patches := gomonkey.ApplyFunc(startGNMIServer, func(_ *TelemetryConfig, _ *gnmi.Config, serverControlSignal <-chan int, wg *sync.WaitGroup) {
defer wg.Done()
})
patches.ApplyFunc(signalHandler, func(serverControlSignal chan<- int, wg *sync.WaitGroup, sigchannel <-chan os.Signal) {
defer wg.Done()
})
defer patches.Reset()

args := []string{"test"}
err := runTelemetry(args)
args := []string{"telemetry", "-logtostderr", "-port", "50051", "-v=2"}
os.Args = args
err := runTelemetry(os.Args)
if err != nil {
t.Errorf("Expected err to be nil, but got %v", err)
}
vflag := flag.Lookup("v")
if vflag.Value.String() != "2" {
t.Errorf("Expected v to be 2")
}
logtostderrflag := flag.Lookup("logtostderr")
if logtostderrflag.Value.String() != "true" {
t.Errorf("Expected logtostderr to be true")
}
}

func TestFlags(t *testing.T) {
Expand Down

0 comments on commit 758ec18

Please sign in to comment.