master
go 181 lines 4.24 KB
Raw
1 package logger
2
3 import (
4 "log/slog"
5 "sync"
6 "sync/atomic"
7 "testing"
8 "time"
9
10 "github.com/stretchr/testify/assert"
11 )
12
13 func TestLimitUsesFixedWindow(t *testing.T) {
14 setTestLevel(t, slog.LevelDebug)
15
16 l, h := newTestLogger(slog.LevelDebug)
17 now := time.Unix(100, 0)
18 l.rl.now = func() time.Time { return now }
19
20 for range 3 {
21 l.Limit("k", 2, 5*time.Second).Info("msg")
22 }
23 assert.Equal(t, 2, h.count())
24
25 now = now.Add(6 * time.Second)
26 l.Limit("k", 2, 5*time.Second).Info("msg")
27 assert.Equal(t, 3, h.count())
28 }
29
30 func TestLimitDZeroIsInfiniteWindow(t *testing.T) {
31 setTestLevel(t, slog.LevelDebug)
32
33 l, h := newTestLogger(slog.LevelDebug)
34 now := time.Unix(200, 0)
35 l.rl.now = func() time.Time { return now }
36
37 l.Limit("k", 2, 0).Info("1")
38 now = now.Add(time.Hour)
39 l.Limit("k", 2, 0).Info("2")
40 now = now.Add(time.Hour)
41 l.Limit("k", 2, 0).Info("3")
42
43 assert.Equal(t, 2, h.count())
44 }
45
46 func TestLimitClampsInvalidInputs(t *testing.T) {
47 setTestLevel(t, slog.LevelDebug)
48
49 l, h := newTestLogger(slog.LevelDebug)
50 l.Limit("k", 0, -time.Second).Info("a")
51 l.Limit("k", 0, -time.Second).Info("b")
52
53 assert.Equal(t, 1, h.count())
54 }
55
56 func TestOnceIsWrapperAndSkipsFormattingWhenSuppressed(t *testing.T) {
57 setTestLevel(t, slog.LevelDebug)
58
59 l, h := newTestLogger(slog.LevelDebug)
60 var hits atomic.Int32
61 p := formatProbe{hits: &hits}
62
63 l.Once("k").Infof("probe %s", p)
64 l.Once("k").Infof("probe %s", p)
65
66 assert.Equal(t, int32(1), hits.Load())
67 assert.Equal(t, 1, h.count())
68 }
69
70 func TestResetAllOnceDoesNotResetLimitState(t *testing.T) {
71 setTestLevel(t, slog.LevelDebug)
72
73 l, h := newTestLogger(slog.LevelDebug)
74 now := time.Unix(300, 0)
75 l.rl.now = func() time.Time { return now }
76
77 l.Once("k").Info("once-1")
78 l.Limit("k", 1, time.Hour).Info("limit-1")
79 l.Once("k").Info("once-2-suppressed")
80 l.Limit("k", 1, time.Hour).Info("limit-2-suppressed")
81
82 assert.Equal(t, 2, h.count())
83
84 l.ResetAllOnce()
85
86 l.Once("k").Info("once-3")
87 l.Limit("k", 1, time.Hour).Info("limit-3-still-suppressed")
88
89 assert.Equal(t, 3, h.count())
90 assert.Equal(t, "once-3", h.last().msg)
91 }
92
93 func TestModeNamespacesAndFirstWriterWinsParams(t *testing.T) {
94 setTestLevel(t, slog.LevelDebug)
95
96 l, h := newTestLogger(slog.LevelDebug)
97 now := time.Unix(400, 0)
98 l.rl.now = func() time.Time { return now }
99
100 // Same key in different modes should be independent.
101 l.Once("shared").Info("once")
102 l.Limit("shared", 1, time.Hour).Info("limit")
103 assert.Equal(t, 2, h.count())
104
105 // First writer wins params for same mode+key.
106 l.Limit("k", 1, time.Hour).Info("first")
107 l.Limit("k", 5, time.Millisecond).Info("second-suppressed")
108 assert.Equal(t, 3, h.count())
109
110 now = now.Add(2 * time.Hour)
111 l.Limit("k", 5, time.Millisecond).Info("third-after-first-window")
112 assert.Equal(t, 4, h.count())
113
114 l.rl.mu.Lock()
115 entry := l.rl.entries[limitKey{mode: modeLimit, key: "k"}]
116 l.rl.mu.Unlock()
117 assert.NotNil(t, entry)
118 assert.Equal(t, 1, entry.limit)
119 assert.Equal(t, time.Hour, entry.window)
120 }
121
122 func TestRateLimiterSweepRemovesStaleEntries(t *testing.T) {
123 setTestLevel(t, slog.LevelDebug)
124
125 l, _ := newTestLogger(slog.LevelDebug)
126 now := time.Unix(500, 0)
127 l.rl.now = func() time.Time { return now }
128 l.rl.ttl = time.Second
129 l.rl.sweepEvery = 1
130
131 l.Once("stale").Info("s")
132 now = now.Add(2 * time.Second)
133 l.Once("fresh").Info("f")
134
135 l.rl.mu.Lock()
136 _, staleOK := l.rl.entries[limitKey{mode: modeOnce, key: "stale"}]
137 _, freshOK := l.rl.entries[limitKey{mode: modeOnce, key: "fresh"}]
138 l.rl.mu.Unlock()
139
140 assert.False(t, staleOK)
141 assert.True(t, freshOK)
142 }
143
144 func TestWithSharesRateLimiterState(t *testing.T) {
145 setTestLevel(t, slog.LevelDebug)
146
147 parent, h := newTestLogger(slog.LevelDebug)
148 child := parent.With("k", "v")
149
150 parent.Once("k").Info("first")
151 child.Once("k").Info("second-suppressed")
152
153 assert.Same(t, parent.rl, child.rl)
154 assert.Equal(t, 1, h.count())
155 }
156
157 func TestOnceConcurrentLogsOnlyOnce(t *testing.T) {
158 setTestLevel(t, slog.LevelDebug)
159
160 l, h := newTestLogger(slog.LevelDebug)
161 var wg sync.WaitGroup
162 for range 100 {
163 wg.Go(func() {
164 l.Once("concurrent").Info("x")
165 })
166 }
167 wg.Wait()
168
169 assert.Equal(t, 1, h.count())
170 }
171
172 func TestRateLimitNilLoggerDoesNotPanic(t *testing.T) {
173 setTestLevel(t, slog.LevelDebug)
174
175 var l *Logger
176 assert.NotPanics(t, func() {
177 l.Once("k").Infof("x=%d", 1)
178 l.Limit("k", 2, time.Second).Warning("y")
179 l.ResetAllOnce()
180 })
181 }