-
Notifications
You must be signed in to change notification settings - Fork 0
/
handler_test.go
162 lines (138 loc) · 3.32 KB
/
handler_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
// Copyright (c) 2024 The sloth authors
// Use of this source code is governed by a MIT license found in the LICENSE file.
package rate_test
import (
"bytes"
"context"
"log/slog"
"runtime"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/nil-go/sloth/internal/assert"
"github.com/nil-go/sloth/rate"
)
func TestNew_panic(t *testing.T) {
t.Parallel()
defer func() {
assert.Equal(t, "cannot create Handler with nil handler", recover().(string))
}()
rate.New(nil)
t.Fail()
}
func TestHandler(t *testing.T) {
t.Parallel()
testcases := []struct {
description string
level slog.Level
expected string
}{
{
description: "level error",
level: slog.LevelError,
expected: `level=ERROR msg=msg pos=first
level=ERROR msg=msg pos=second
level=ERROR msg=msg pos=fourth
level=ERROR msg=msg g.pos=after
`,
},
{
description: "level warn",
level: slog.LevelWarn,
expected: `level=WARN msg=msg pos=first
level=WARN msg=msg pos=second
level=WARN msg=msg pos=fourth
level=WARN msg=msg g.pos=after
`,
},
{
description: "level info",
level: slog.LevelInfo,
expected: `level=INFO msg=msg pos=first
level=INFO msg=msg pos=second
level=INFO msg=msg pos=fourth
level=INFO msg=msg g.pos=after
`,
},
{
description: "level debug",
level: slog.LevelDebug,
expected: `level=DEBUG msg=msg pos=first
level=DEBUG msg=msg pos=second
level=DEBUG msg=msg pos=fourth
level=DEBUG msg=msg g.pos=after
`,
},
}
for _, testcase := range testcases {
t.Run(testcase.description, func(t *testing.T) {
t.Parallel()
buf := &bytes.Buffer{}
handler := rate.New(
slog.NewTextHandler(buf, &slog.HandlerOptions{
Level: slog.LevelDebug,
ReplaceAttr: func(groups []string, attr slog.Attr) slog.Attr {
if len(groups) == 0 && attr.Key == slog.TimeKey {
return slog.Attr{}
}
return attr
},
}),
rate.WithFirst(2),
rate.WithEvery(2),
rate.WithInterval(time.Second),
)
logger := slog.New(handler)
ctx := context.Background()
logger.Log(ctx, testcase.level, "msg", "pos", "first")
logger.Log(ctx, testcase.level, "msg", "pos", "second")
logger.Log(ctx, testcase.level, "msg", "pos", "third")
logger.Log(ctx, testcase.level, "msg", "pos", "fourth")
time.Sleep(time.Second)
logger.WithGroup("g").With("pos", "after").Log(ctx, testcase.level, "msg")
assert.Equal(t, testcase.expected, buf.String())
})
}
}
func TestHandler_race(t *testing.T) {
t.Parallel()
procs := runtime.GOMAXPROCS(0)
counter := atomic.Int64{}
handler := rate.New(
countHandler{count: &counter},
rate.WithFirst(1),
rate.WithEvery(1000),
)
logger := slog.New(handler)
ctx := context.Background()
start := make(chan struct{})
var waitGroup sync.WaitGroup
waitGroup.Add(procs)
for range procs {
go func() {
defer waitGroup.Done()
<-start
logger.Log(ctx, slog.LevelInfo, "msg")
}()
}
close(start)
waitGroup.Wait()
assert.Equal(t, 1, int(counter.Load()))
}
type countHandler struct {
count *atomic.Int64
}
func (c countHandler) Enabled(context.Context, slog.Level) bool {
return true
}
func (c countHandler) Handle(context.Context, slog.Record) error {
c.count.Add(1)
return nil
}
func (c countHandler) WithAttrs([]slog.Attr) slog.Handler {
return c
}
func (c countHandler) WithGroup(string) slog.Handler {
return c
}