starlog/rate_limit_test.go

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)
	}
}