package starlog import ( "bytes" "strings" "sync" "testing" "time" ) func setRateLimiterNowForTest(logger *StarLogger, nowFunc func() time.Time) { if logger == nil || logger.logcore == nil { return } logger.logcore.mu.Lock() limiter := logger.logcore.rateLimiter logger.logcore.mu.Unlock() if limiter == nil { return } limiter.setNowFuncForTest(nowFunc) } func TestRateLimitDrop(t *testing.T) { var buf bytes.Buffer logger := newStructuredTestLogger(&buf) dropped := make(chan RateLimitDropData, 1) cfg := DefaultRateLimitConfig() cfg.Enable = true cfg.Levels = []int{LvInfo} cfg.Rate = 1 cfg.Burst = 1 cfg.DropPolicy = RateLimitDrop cfg.OnDrop = func(data RateLimitDropData) { select { case dropped <- data: default: } } logger.SetRateLimitConfig(cfg) now := time.Date(2026, 3, 19, 10, 0, 0, 0, time.UTC) setRateLimiterNowForTest(logger, func() time.Time { return now }) logger.Info("first") logger.Info("second") got := buf.String() if !strings.Contains(got, "first") { t.Fatalf("first log should pass, got %q", got) } if strings.Contains(got, "second") { t.Fatalf("second log should be dropped, got %q", got) } select { case data := <-dropped: if data.PassThrough { t.Fatalf("drop policy should not pass through, got %+v", data) } if data.Key == "" { t.Fatalf("drop callback should provide key") } default: t.Fatalf("drop callback should be invoked") } stats := logger.GetRateLimitStats() if stats.Allowed != 1 || stats.Dropped != 1 || stats.Suppressed != 1 { t.Fatalf("unexpected rate limit stats: %+v", stats) } } func TestRateLimitPassThrough(t *testing.T) { var buf bytes.Buffer logger := newStructuredTestLogger(&buf) dropped := make(chan RateLimitDropData, 1) cfg := DefaultRateLimitConfig() cfg.Enable = true cfg.Levels = []int{LvInfo} cfg.Rate = 1 cfg.Burst = 1 cfg.DropPolicy = RateLimitPassThrough cfg.OnDrop = func(data RateLimitDropData) { select { case dropped <- data: default: } } logger.SetRateLimitConfig(cfg) now := time.Date(2026, 3, 19, 10, 1, 0, 0, time.UTC) setRateLimiterNowForTest(logger, func() time.Time { return now }) logger.Info("one") logger.Info("two") got := buf.String() if !strings.Contains(got, "one") || !strings.Contains(got, "two") { t.Fatalf("pass-through mode should keep both logs, got %q", got) } select { case data := <-dropped: if !data.PassThrough { t.Fatalf("drop callback should mark pass-through mode, got %+v", data) } default: t.Fatalf("drop callback should be invoked in pass-through mode") } stats := logger.GetRateLimitStats() if stats.Allowed != 2 || stats.Dropped != 0 || stats.PassedThrough != 1 || stats.Suppressed != 1 { t.Fatalf("unexpected pass-through stats: %+v", stats) } } func TestRateLimitLevelsAndExempt(t *testing.T) { var buf bytes.Buffer logger := newStructuredTestLogger(&buf) cfg := DefaultRateLimitConfig() cfg.Enable = true cfg.Levels = []int{LvInfo} cfg.ExemptLevels = []int{LvError} cfg.Rate = 1 cfg.Burst = 1 logger.SetRateLimitConfig(cfg) now := time.Date(2026, 3, 19, 10, 2, 0, 0, time.UTC) setRateLimiterNowForTest(logger, func() time.Time { return now }) logger.Debug("d1") logger.Debug("d2") logger.Info("i1") logger.Info("i2") logger.Error("e1") logger.Error("e2") got := buf.String() if !strings.Contains(got, "d1") || !strings.Contains(got, "d2") { t.Fatalf("debug logs should not be limited, got %q", got) } if !strings.Contains(got, "i1") || strings.Contains(got, "i2") { t.Fatalf("info logs should be limited, got %q", got) } if !strings.Contains(got, "e1") || !strings.Contains(got, "e2") { t.Fatalf("error logs should be exempt, got %q", got) } stats := logger.GetRateLimitStats() if stats.Dropped != 1 || stats.Allowed != 5 { t.Fatalf("unexpected level/exempt stats: %+v", stats) } } func TestRateLimitByKey(t *testing.T) { var buf bytes.Buffer logger := newStructuredTestLogger(&buf) cfg := DefaultRateLimitConfig() cfg.Enable = true cfg.Levels = []int{LvInfo} cfg.Scope = RateLimitScopeByKey cfg.Rate = 1 cfg.Burst = 1 cfg.MaxKeys = 16 cfg.KeyFunc = func(entry *Entry) string { if entry == nil || entry.Fields == nil { return "" } if user, ok := entry.Fields["user"].(string); ok { return user } return "" } logger.SetRateLimitConfig(cfg) now := time.Date(2026, 3, 19, 10, 3, 0, 0, time.UTC) setRateLimiterNowForTest(logger, func() time.Time { return now }) logger.WithField("user", "a").Info("key-a-one") logger.WithField("user", "a").Info("key-a-two") logger.WithField("user", "b").Info("key-b-one") got := buf.String() if !strings.Contains(got, "key-a-one") || strings.Contains(got, "key-a-two") { t.Fatalf("same key should be limited, got %q", got) } if !strings.Contains(got, "key-b-one") { t.Fatalf("different key should have independent bucket, got %q", got) } stats := logger.GetRateLimitStats() if stats.Dropped != 1 || stats.Allowed != 2 || stats.CurrentKeys < 2 { t.Fatalf("unexpected by-key stats: %+v", stats) } } func TestRateLimitSummary(t *testing.T) { var buf bytes.Buffer logger := newStructuredTestLogger(&buf) var mu sync.Mutex callbacks := make([]RateLimitDropData, 0, 4) cfg := DefaultRateLimitConfig() cfg.Enable = true cfg.Levels = []int{LvInfo} cfg.Rate = 0.001 cfg.Burst = 1 cfg.SummaryInterval = 2 * time.Second cfg.OnDrop = func(data RateLimitDropData) { mu.Lock() callbacks = append(callbacks, data) mu.Unlock() } logger.SetRateLimitConfig(cfg) now := time.Date(2026, 3, 19, 10, 4, 0, 0, time.UTC) setRateLimiterNowForTest(logger, func() time.Time { return now }) logger.Info("one") logger.Info("two") now = now.Add(3 * time.Second) logger.Info("three") mu.Lock() defer mu.Unlock() if len(callbacks) != 2 { t.Fatalf("expected 2 drop callbacks, got %d", len(callbacks)) } if callbacks[0].Summary { t.Fatalf("first drop should not be summary, got %+v", callbacks[0]) } if !callbacks[1].Summary { t.Fatalf("second drop should trigger summary, got %+v", callbacks[1]) } if callbacks[1].SummarySuppressed != 2 { t.Fatalf("summary should include suppressed count, got %+v", callbacks[1]) } stats := logger.GetRateLimitStats() if stats.Summaries != 1 { t.Fatalf("expected one summary event, got %+v", stats) } } func TestRateLimitConfigSnapshotApply(t *testing.T) { logger := NewStarlog(nil) cfg := DefaultRateLimitConfig() cfg.Enable = true cfg.Levels = []int{LvWarning, LvError} cfg.Rate = 5 cfg.Burst = 10 cfg.Scope = RateLimitScopeByKey cfg.MaxKeys = 99 cfg.KeyTTL = time.Minute cfg.DropPolicy = RateLimitPassThrough logger.SetRateLimitConfig(cfg) snapshot := logger.GetConfig() if !snapshot.RateLimit.Enable || snapshot.RateLimit.Rate != 5 || snapshot.RateLimit.Burst != 10 { t.Fatalf("config snapshot should include rate limit config, got %+v", snapshot.RateLimit) } if snapshot.RateLimit.Scope != RateLimitScopeByKey || snapshot.RateLimit.DropPolicy != RateLimitPassThrough { t.Fatalf("unexpected rate limit snapshot details: %+v", snapshot.RateLimit) } snapshot.RateLimit.Enable = false snapshot.RateLimit.Rate = 1 snapshot.RateLimit.Burst = 1 logger.ApplyConfig(snapshot) current := logger.GetRateLimitConfig() if current.Enable { t.Fatalf("apply config should update rate limit enable flag") } if current.Rate != 1 || current.Burst != 1 { t.Fatalf("apply config should update rate limit values, got %+v", current) } } func TestStdRateLimitBridge(t *testing.T) { backup := GetConfig() defer ApplyConfig(backup) cfg := DefaultRateLimitConfig() cfg.Enable = true cfg.Rate = 2 cfg.Burst = 2 SetRateLimitConfig(cfg) got := GetRateLimitConfig() if !got.Enable || got.Rate != 2 || got.Burst != 2 { t.Fatalf("std bridge should set/get rate limit config, got %+v", got) } EnableRateLimit(false) got = GetRateLimitConfig() if got.Enable { t.Fatalf("EnableRateLimit(false) should disable limiter") } ResetRateLimitStats() stats := GetRateLimitStats() if stats.Allowed != 0 || stats.Dropped != 0 || stats.Suppressed != 0 { t.Fatalf("reset stats should clear counters, got %+v", stats) } }