Skip to content

Commit

Permalink
Make single point of exit
Browse files Browse the repository at this point in the history
  • Loading branch information
@gandarez
gandarez committed Jul 23, 2024
1 parent d9c9299 commit 1415ea2
Show file tree
Hide file tree
Showing 5 changed files with 343 additions and 447 deletions.
19 changes: 17 additions & 2 deletions cmd/root.go
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
package cmd

import (
"errors"
"fmt"
"os"

"github.com/wakatime/wakatime-cli/pkg/api"
"github.com/wakatime/wakatime-cli/pkg/exitcode"
"github.com/wakatime/wakatime-cli/pkg/offline"

log "github.com/sirupsen/logrus"
Expand All @@ -25,8 +28,20 @@ func NewRootCMD() *cobra.Command {
cmd := &cobra.Command{
Use: "wakatime-cli",
Short: "Command line interface used by all WakaTime text editor plugins.",
Run: func(cmd *cobra.Command, _ []string) {
Run(cmd, v)
RunE: func(cmd *cobra.Command, _ []string) error {
if err := RunE(cmd, v); err != nil {
var errexitcode exitcode.Err

if errors.As(err, &errexitcode) {
os.Exit(errexitcode.Code)
}

os.Exit(exitcode.ErrGeneric)
}

os.Exit(exitcode.Success)

return nil
},
}

Expand Down
97 changes: 45 additions & 52 deletions cmd/run.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,8 @@ type diagnostics struct {
Stack string
}

// Run executes commands parsed from a command line.
func Run(cmd *cobra.Command, v *viper.Viper) {
// RunE executes commands parsed from a command line.
func RunE(cmd *cobra.Command, v *viper.Viper) error {
// force setup logging otherwise log goes to std out
_, err := SetupLogging(v)
if err != nil {
Expand All @@ -58,9 +58,9 @@ func Run(cmd *cobra.Command, v *viper.Viper) {
log.Errorf("failed to parse config files: %s", err)

if v.IsSet("entity") {
saveHeartbeats(v)
_ = saveHeartbeats(v)

os.Exit(exitcode.ErrConfigFileParse)
return exitcode.Err{Code: exitcode.ErrConfigFileParse}
}
}

Expand All @@ -75,90 +75,88 @@ func Run(cmd *cobra.Command, v *viper.Viper) {
log.Fatalf("failed to register custom lexers: %s", err)
}

shutdown := shutdownFn(func() {})

// start profiling if enabled
if logFileParams.Metrics {
shutdown, err = metrics.StartProfiling()
shutdown, err := metrics.StartProfiling()
if err != nil {
log.Errorf("failed to start profiling: %s", err)
}

defer shutdown()
}

if v.GetBool("user-agent") {
log.Debugln("command: user-agent")

fmt.Println(heartbeat.UserAgent(vipertools.GetString(v, "plugin")))

shutdown()

os.Exit(exitcode.Success)
return nil
}

if v.GetBool("version") {
log.Debugln("command: version")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, runVersion, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, runVersion)
}

if v.IsSet("config-read") {
log.Debugln("command: config-read")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, configread.Run, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, configread.Run)
}

if v.IsSet("config-write") {
log.Debugln("command: config-write")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, configwrite.Run, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, configwrite.Run)
}

if v.GetBool("today") {
log.Debugln("command: today")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, today.Run, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, today.Run)
}

if v.IsSet("today-goal") {
log.Debugln("command: today-goal")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, todaygoal.Run, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, todaygoal.Run)
}

if v.GetBool("file-experts") {
log.Debugln("command: file-experts")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, fileexperts.Run, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, fileexperts.Run)
}

if v.IsSet("entity") {
log.Debugln("command: heartbeat")

if v.GetBool("offline-only") {
saveHeartbeats(v)
shutdown()
os.Exit(exitcode.Success)
exitCode := saveHeartbeats(v)

os.Exit(exitCode) // nolint:gocritic
}

RunCmdWithOfflineSync(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, cmdheartbeat.Run, shutdown)
return RunCmdWithOfflineSync(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, cmdheartbeat.Run)
}

if v.IsSet("sync-offline-activity") {
log.Debugln("command: sync-offline-activity")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlinesync.Run, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlinesync.Run)
}

if v.GetBool("offline-count") {
log.Debugln("command: offline-count")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlinecount.Run, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlinecount.Run)
}

if v.IsSet("print-offline-heartbeats") {
log.Debugln("command: print-offline-heartbeats")

RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlineprint.Run, shutdown)
return RunCmd(v, logFileParams.Verbose, logFileParams.SendDiagsOnErrors, offlineprint.Run)
}

log.Warnf("one of the following parameters has to be provided: %s", strings.Join([]string{
Expand All @@ -177,7 +175,7 @@ func Run(cmd *cobra.Command, v *viper.Viper) {

_ = cmd.Help()

os.Exit(exitcode.ErrGeneric)
return exitcode.Err{Code: exitcode.ErrGeneric}
}

func parseConfigFiles(v *viper.Viper) error {
Expand Down Expand Up @@ -264,46 +262,31 @@ func SetupLogging(v *viper.Viper) (*logfile.Params, error) {
return &logfileParams, nil
}

type (
// cmdFn represents a command function.
cmdFn func(v *viper.Viper) (int, error)
// shutdownFn represents a shutdown function. It will be called before exiting.
shutdownFn func()
)
// cmdFn represents a command function.
type cmdFn func(v *viper.Viper) (int, error)

// RunCmd runs a command function and exits with the exit code returned by
// the command function. Will send diagnostic on any errors or panics.
func RunCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn, shutdown shutdownFn) {
exitCode := runCmd(v, verbose, sendDiagsOnErrors, cmd)

shutdown()

os.Exit(exitCode)
func RunCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) error {
return runCmd(v, verbose, sendDiagsOnErrors, cmd)
}

// RunCmdWithOfflineSync runs a command function and exits with the exit code
// returned by the command function. If command run was successful, it will execute
// offline sync command afterwards. Will send diagnostic on any errors or panics.
func RunCmdWithOfflineSync(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn, shutdown shutdownFn) {
exitCode := runCmd(v, verbose, sendDiagsOnErrors, cmd)
if exitCode != exitcode.Success {
shutdown()

os.Exit(exitCode)
func RunCmdWithOfflineSync(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) error {
if err := runCmd(v, verbose, sendDiagsOnErrors, cmd); err != nil {
return err
}

exitCode = runCmd(v, verbose, sendDiagsOnErrors, offlinesync.Run)

shutdown()

os.Exit(exitCode)
return runCmd(v, verbose, sendDiagsOnErrors, offlinesync.Run)
}

// runCmd contains the main logic of RunCmd.
// It will send diagnostic on any errors or panics.
// On panic, it will send diagnostic and exit with ErrGeneric exit code.
// On error, it will only send diagnostic if sendDiagsOnErrors and verbose is true.
func runCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) (exitCode int) {
func runCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) (errresponse error) {
logs := bytes.NewBuffer(nil)
resetLogs := captureLogs(logs)

Expand All @@ -328,14 +311,14 @@ func runCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) (ex
log.Warnf("failed to send diagnostics: %s", err)
}

exitCode = exitcode.ErrGeneric
errresponse = exitcode.Err{Code: exitcode.ErrGeneric}
}
}()

var err error

// run command
exitCode, err = cmd(v)
exitCode, err := cmd(v)
// nolint:nestif
if err != nil {
if errwaka, ok := err.(wakaerror.Error); ok {
Expand All @@ -362,18 +345,28 @@ func runCmd(v *viper.Viper, verbose bool, sendDiagsOnErrors bool, cmd cmdFn) (ex
}
}

return exitCode
if exitCode != exitcode.Success {
log.Debugf("command failed with exit code %d", exitCode)

errresponse = exitcode.Err{Code: exitCode}
}

return errresponse
}

func saveHeartbeats(v *viper.Viper) {
func saveHeartbeats(v *viper.Viper) int {
queueFilepath, err := offline.QueueFilepath()
if err != nil {
log.Warnf("failed to load offline queue filepath: %s", err)
}

if err := cmdoffline.SaveHeartbeats(v, nil, queueFilepath); err != nil {
log.Errorf("failed to save heartbeats to offline queue: %s", err)

return exitcode.ErrGeneric
}

return exitcode.Success
}

func sendDiagnostics(v *viper.Viper, d diagnostics) error {
Expand Down
40 changes: 28 additions & 12 deletions cmd/run_internal_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,21 +25,25 @@ import (
func TestRunCmd(t *testing.T) {
v := viper.New()

ret := runCmd(v, false, false, func(_ *viper.Viper) (int, error) {
err := runCmd(v, false, false, func(_ *viper.Viper) (int, error) {
return exitcode.Success, nil
})

assert.Equal(t, exitcode.Success, ret)
assert.Nil(t, err)
}

func TestRunCmd_Err(t *testing.T) {
v := viper.New()

ret := runCmd(v, false, false, func(_ *viper.Viper) (int, error) {
err := runCmd(v, false, false, func(_ *viper.Viper) (int, error) {
return exitcode.ErrGeneric, errors.New("fail")
})

assert.Equal(t, exitcode.ErrGeneric, ret)
var errexitcode exitcode.Err

require.ErrorAs(t, err, &errexitcode)

assert.Equal(t, exitcode.ErrGeneric, err.(exitcode.Err).Code)
}

func TestRunCmd_ErrOfflineEnqueue(t *testing.T) {
Expand Down Expand Up @@ -95,11 +99,15 @@ func TestRunCmd_ErrOfflineEnqueue(t *testing.T) {
v.Set("key", "00000000-0000-4000-8000-000000000000")
v.Set("plugin", "vim")

ret := runCmd(v, true, false, func(_ *viper.Viper) (int, error) {
err := runCmd(v, true, false, func(_ *viper.Viper) (int, error) {
return exitcode.ErrGeneric, errors.New("fail")
})

assert.Equal(t, exitcode.ErrGeneric, ret)
var errexitcode exitcode.Err

require.ErrorAs(t, err, &errexitcode)

assert.Equal(t, exitcode.ErrGeneric, err.(exitcode.Err).Code)
}

func TestRunCmd_BackoffLoggedWithVerbose(t *testing.T) {
Expand Down Expand Up @@ -147,10 +155,15 @@ func TestRunCmd_BackoffLoggedWithVerbose(t *testing.T) {
v.Set("internal.backoff_retries", "1")
v.Set("verbose", verbose)

SetupLogging(v)
_, _ = SetupLogging(v)

exitCode := runCmd(v, verbose, false, cmdheartbeat.Run)
assert.Equal(t, exitcode.ErrBackoff, exitCode)
err = runCmd(v, verbose, false, cmdheartbeat.Run)

var errexitcode exitcode.Err

require.ErrorAs(t, err, &errexitcode)

assert.Equal(t, exitcode.ErrBackoff, err.(exitcode.Err).Code)

assert.Equal(t, 0, numCalls)

Expand Down Expand Up @@ -205,11 +218,14 @@ func TestRunCmd_BackoffNotLogged(t *testing.T) {
v.Set("internal.backoff_retries", "1")
v.Set("verbose", verbose)

SetupLogging(v)
_, _ = SetupLogging(v)

err = runCmd(v, verbose, false, cmdheartbeat.Run)

exitCode := runCmd(v, verbose, false, cmdheartbeat.Run)
assert.Equal(t, exitcode.ErrBackoff, exitCode)
var errexitcode exitcode.Err

require.ErrorAs(t, err, &errexitcode)
assert.Equal(t, exitcode.ErrBackoff, err.(exitcode.Err).Code)
assert.Equal(t, 0, numCalls)

output, err := io.ReadAll(logFile)
Expand Down
Loading

0 comments on commit 1415ea2

Please sign in to comment.