Fix rate limit custom number

wakatime · Jan 6, 2025 · ce559c9 · ce559c9
1 parent 92addb5
commit ce559c9
Show file tree

Hide file tree

Showing 8 changed files with 46 additions and 19 deletions.
diff --git a/cmd/params/params.go b/cmd/params/params.go
@@ -22,6 +22,7 @@ import (
 	"github.com/wakatime/wakatime-cli/pkg/heartbeat"
 	"github.com/wakatime/wakatime-cli/pkg/ini"
 	"github.com/wakatime/wakatime-cli/pkg/log"
+	"github.com/wakatime/wakatime-cli/pkg/offline"
 	"github.com/wakatime/wakatime-cli/pkg/output"
 	"github.com/wakatime/wakatime-cli/pkg/project"
 	"github.com/wakatime/wakatime-cli/pkg/regex"
@@ -658,17 +659,26 @@ func LoadOfflineParams(ctx context.Context, v *viper.Viper) Offline {
 
 	logger := log.Extract(ctx)
 
-	rateLimit, _ := vipertools.FirstNonEmptyInt(v, "heartbeat-rate-limit-seconds", "settings.heartbeat_rate_limit_seconds")
-	if rateLimit < 0 {
-		logger.Warnf("argument --heartbeat-rate-limit-seconds must be zero or a positive integer number, got %d", rateLimit)
+	rateLimit := offline.RateLimitDefaultSeconds
 
-		rateLimit = 0
+	if rateLimitSecs, ok := vipertools.FirstNonEmptyInt(v,
+		"heartbeat-rate-limit-seconds",
+		"settings.heartbeat_rate_limit_seconds"); ok {
+		rateLimit = rateLimitSecs
+
+		if rateLimit < 0 {
+			logger.Warnf(
+				"argument --heartbeat-rate-limit-seconds must be zero or a positive integer number, got %d",
+				rateLimit,
+			)
+
+			rateLimit = 0
+		}
 	}
 
 	syncMax := v.GetInt("sync-offline-activity")
 	if syncMax < 0 {
 		logger.Warnf("argument --sync-offline-activity must be zero or a positive integer number, got %d", syncMax)
-
 		syncMax = 0
 	}
 
@@ -1100,7 +1110,7 @@ func (p Offline) String() string {
 	}
 
 	return fmt.Sprintf(
-		"disabled: %t, last sent at: '%s', print max: %d, num rate limit: %d, num sync max: %d",
+		"disabled: %t, last sent at: '%s', print max: %d, rate limit: %s, num sync max: %d",
 		p.Disabled,
 		lastSentAt,
 		p.PrintMax,

diff --git a/cmd/params/params_test.go b/cmd/params/params_test.go
@@ -20,14 +20,15 @@ import (
 	"github.com/wakatime/wakatime-cli/pkg/heartbeat"
 	inipkg "github.com/wakatime/wakatime-cli/pkg/ini"
 	"github.com/wakatime/wakatime-cli/pkg/log"
+	"github.com/wakatime/wakatime-cli/pkg/offline"
 	"github.com/wakatime/wakatime-cli/pkg/output"
 	"github.com/wakatime/wakatime-cli/pkg/project"
 	"github.com/wakatime/wakatime-cli/pkg/regex"
-	"gopkg.in/ini.v1"
 
 	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gopkg.in/ini.v1"
 )
 
 func TestLoadHeartbeatParams_AlternateProject(t *testing.T) {
@@ -1832,7 +1833,7 @@ func TestLoadOfflineParams_RateLimit_FromConfig(t *testing.T) {
 
 func TestLoadOfflineParams_RateLimit_Zero(t *testing.T) {
 	v := viper.New()
-	v.Set("heartbeat-rate-limit-seconds", "0")
+	v.Set("heartbeat-rate-limit-seconds", 0)
 
 	params := cmdparams.LoadOfflineParams(context.Background(), v)
 
@@ -1841,11 +1842,11 @@ func TestLoadOfflineParams_RateLimit_Zero(t *testing.T) {
 
 func TestLoadOfflineParams_RateLimit_Default(t *testing.T) {
 	v := viper.New()
-	v.SetDefault("heartbeat-rate-limit-seconds", 20)
+	v.SetDefault("heartbeat-rate-limit-seconds", offline.RateLimitDefaultSeconds)
 
 	params := cmdparams.LoadOfflineParams(context.Background(), v)
 
-	assert.Equal(t, time.Duration(20)*time.Second, params.RateLimit)
+	assert.Equal(t, time.Duration(offline.RateLimitDefaultSeconds)*time.Second, params.RateLimit)
 }
 
 func TestLoadOfflineParams_RateLimit_NegativeNumber(t *testing.T) {
@@ -1863,7 +1864,7 @@ func TestLoadOfflineParams_RateLimit_NonIntegerValue(t *testing.T) {
 
 	params := cmdparams.LoadOfflineParams(context.Background(), v)
 
-	assert.Zero(t, params.RateLimit)
+	assert.Equal(t, time.Duration(offline.RateLimitDefaultSeconds)*time.Second, params.RateLimit)
 }
 
 func TestLoadOfflineParams_LastSentAt(t *testing.T) {
@@ -1889,7 +1890,7 @@ func TestLoadOfflineParams_LastSentAt_Err(t *testing.T) {
 
 func TestLoadOfflineParams_LastSentAtFuture(t *testing.T) {
 	v := viper.New()
-	lastSentAt := time.Now().Add(time.Duration(2) * time.Hour)
+	lastSentAt := time.Now().Add(2 * time.Hour)
 	v.Set("internal.heartbeats_last_sent_at", lastSentAt.Format(inipkg.DateFormat))
 
 	params := cmdparams.LoadOfflineParams(context.Background(), v)
@@ -2658,14 +2659,14 @@ func TestOffline_String(t *testing.T) {
 		Disabled:   true,
 		LastSentAt: lastSentAt,
 		PrintMax:   6,
-		RateLimit:  15,
+		RateLimit:  time.Duration(15) * time.Second,
 		SyncMax:    12,
 	}
 
 	assert.Equal(
 		t,
 		"disabled: true, last sent at: '2021-08-30T18:50:42-03:00', print max: 6,"+
-			" num rate limit: 15, num sync max: 12",
+			" rate limit: 15s, num sync max: 12",
 		offline.String(),
 	)
 }

diff --git a/cmd/root.go b/cmd/root.go
@@ -128,7 +128,7 @@ func setFlags(cmd *cobra.Command, v *viper.Viper) {
 		"heartbeat-rate-limit-seconds",
 		offline.RateLimitDefaultSeconds,
 		fmt.Sprintf("Only sync heartbeats to the API once per these seconds, instead"+
-			" saving to the offline db. Defaults to %d. Use zero to disable.",
+			" saving to the offline db. Defaults to %d. Use 0 to disable.",
 			offline.RateLimitDefaultSeconds),
 	)
 	flags.String("hide-branch-names", "", "Obfuscate branch names. Will not send revision control branch names to api.")

diff --git a/cmd/run_internal_test.go b/cmd/run_internal_test.go
@@ -415,6 +415,10 @@ func TestParseConfigFiles(t *testing.T) {
 	assert.Equal(t,
 		"2006-01-02T15:04:05Z07:00",
 		v.GetString("internal.backoff_at"))
+	assert.Equal(t,
+		"2025-01-05T22:21:51Z03:00",
+		v.GetString("internal.heartbeats_last_sent_at"),
+	)
 }
 
 func TestParseConfigFiles_MissingAPIKey(t *testing.T) {

diff --git a/cmd/testdata/.wakatime-internal.cfg b/cmd/testdata/.wakatime-internal.cfg
@@ -1,3 +1,4 @@
 [internal]
 backoff_retries  = 1
 backoff_at       = 2006-01-02T15:04:05Z07:00
+heartbeats_last_sent_at = 2025-01-05T22:21:51Z03:00
diff --git a/main_test.go b/main_test.go
@@ -455,7 +455,6 @@ func TestSendHeartbeats_ExtraHeartbeats_SyncLegacyOfflineActivity(t *testing.T)
 		"--offline-queue-file-legacy", offlineQueueFileLegacy.Name(),
 		"--lineno", "42",
 		"--lines-in-file", "100",
-		"--heartbeat-rate-limit-seconds", "0",
 		"--time", "1585598059",
 		"--hide-branch-names", ".*",
 		"--write",

diff --git a/pkg/vipertools/vipertools.go b/pkg/vipertools/vipertools.go
@@ -3,6 +3,7 @@ package vipertools
 import (
 	"strings"
 
+	"github.com/spf13/cast"
 	"github.com/spf13/viper"
 )
 
@@ -31,9 +32,20 @@ func FirstNonEmptyInt(v *viper.Viper, keys ...string) (int, bool) {
 	}
 
 	for _, key := range keys {
-		if value := v.GetInt(key); value != 0 {
-			return value, true
+		if !v.IsSet(key) {
+			continue
+		}
+
+		// Zero means a valid value when set, so it needs to use generic function and later cast it to int
+		value := v.Get(key)
+
+		// If the value is not an int, it will continue to find the next non-empty key
+		parsed, err := cast.ToIntE(value)
+		if err != nil {
+			continue
 		}
+
+		return parsed, true
 	}
 
 	return 0, false

diff --git a/pkg/vipertools/vipertools_test.go b/pkg/vipertools/vipertools_test.go
@@ -66,7 +66,7 @@ func TestFirstNonEmptyInt_EmptyInt(t *testing.T) {
 	v := viper.New()
 	v.Set("first", 0)
 	_, ok := vipertools.FirstNonEmptyInt(v, "first")
-	assert.False(t, ok)
+	assert.True(t, ok)
 }
 
 func TestFirstNonEmptyInt_StringValue(t *testing.T) {