capturing.go 4.58 KB
Newer Older
1 2 3
package testlog

import (
4
	"context"
5
	"log/slog"
6 7
	"strings"

8 9 10
	"github.com/ethereum/go-ethereum/log"
)

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
// CapturedAttributes forms a chain of inherited attributes, to traverse on captured log records.
type CapturedAttributes struct {
	Parent     *CapturedAttributes
	Attributes []slog.Attr
}

// Attrs calls f on each Attr in the [CapturedAttributes].
// Iteration stops if f returns false.
func (r *CapturedAttributes) Attrs(f func(slog.Attr) bool) {
	for _, a := range r.Attributes {
		if !f(a) {
			return
		}
	}
	if r.Parent != nil {
		r.Parent.Attrs(f)
	}
}

// CapturedRecord is a wrapped around a regular log-record,
// to preserve the inherited attributes context, without mutating the record or reordering attributes.
type CapturedRecord struct {
	Parent *CapturedAttributes
	*slog.Record
}

// Attrs calls f on each Attr in the [CapturedRecord].
// Iteration stops if f returns false.
func (r *CapturedRecord) Attrs(f func(slog.Attr) bool) {
	searching := true
	r.Record.Attrs(func(a slog.Attr) bool {
		searching = f(a)
		return searching
	})
	if !searching { // if we found it already, then don't traverse the remainder
		return
	}
	if r.Parent != nil {
		r.Parent.Attrs(f)
	}
}

53 54 55
// CapturingHandler provides a log handler that captures all log records and optionally forwards them to a delegate.
// Note that it is not thread safe.
type CapturingHandler struct {
56
	handler slog.Handler
57 58 59
	Logs    *[]*CapturedRecord // shared among derived CapturingHandlers
	// attrs are inherited log record attributes, from a logger that this CapturingHandler may be derived from
	attrs *CapturedAttributes
60 61
}

62 63
func CaptureLogger(t Testing, level slog.Level) (_ log.Logger, ch *CapturingHandler) {
	return LoggerWithHandlerMod(t, level, func(h slog.Handler) slog.Handler {
64
		ch = &CapturingHandler{handler: h, Logs: new([]*CapturedRecord)}
65 66 67 68 69 70 71 72
		return ch
	}), ch
}

func (c *CapturingHandler) Enabled(context.Context, slog.Level) bool {
	// We want to capture all logs, even if the underlying handler only logs
	// above a certain level.
	return true
73 74
}

75
func (c *CapturingHandler) Handle(ctx context.Context, r slog.Record) error {
76 77 78 79
	*c.Logs = append(*c.Logs, &CapturedRecord{
		Parent: c.attrs,
		Record: &r,
	})
80 81
	if c.handler != nil && c.handler.Enabled(ctx, r.Level) {
		return c.handler.Handle(ctx, r)
82 83 84 85
	}
	return nil
}

86 87 88 89
func (c *CapturingHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return &CapturingHandler{
		handler: c.handler.WithAttrs(attrs),
		Logs:    c.Logs,
90 91 92 93
		attrs: &CapturedAttributes{
			Parent:     c.attrs,
			Attributes: attrs,
		},
94 95 96 97 98 99 100 101 102 103
	}
}

func (c *CapturingHandler) WithGroup(name string) slog.Handler {
	return &CapturingHandler{
		handler: c.handler.WithGroup(name),
		Logs:    c.Logs,
	}
}

104
func (c *CapturingHandler) Clear() {
105
	*c.Logs = (*c.Logs)[:0] // reuse slice
106 107
}

108
func NewLevelFilter(level slog.Level) LogFilter {
109 110
	return func(r *CapturedRecord) bool {
		return r.Record.Level == level
111 112 113
	}
}

114
func NewAttributesFilter(key, value string) LogFilter {
115
	return func(r *CapturedRecord) bool {
116 117 118 119 120 121 122 123 124 125 126 127 128
		found := false
		r.Attrs(func(a slog.Attr) bool {
			if a.Key == key && a.Value.String() == value {
				found = true
				return false
			}
			return true // try next
		})
		return found
	}
}

func NewAttributesContainsFilter(key, value string) LogFilter {
129
	return func(r *CapturedRecord) bool {
130 131 132 133 134 135 136 137 138 139 140 141
		found := false
		r.Attrs(func(a slog.Attr) bool {
			if a.Key == key && strings.Contains(a.Value.String(), value) {
				found = true
				return false
			}
			return true // try next
		})
		return found
	}
}

142
func NewMessageFilter(message string) LogFilter {
143 144
	return func(r *CapturedRecord) bool {
		return r.Record.Message == message
145 146 147 148
	}
}

func NewMessageContainsFilter(message string) LogFilter {
149 150
	return func(r *CapturedRecord) bool {
		return strings.Contains(r.Record.Message, message)
151 152 153
	}
}

154
type LogFilter func(record *CapturedRecord) bool
155

156
func (c *CapturingHandler) FindLog(filters ...LogFilter) *CapturedRecord {
157
	for _, record := range *c.Logs {
158 159 160 161 162 163 164 165
		match := true
		for _, filter := range filters {
			if !filter(record) {
				match = false
				break
			}
		}
		if match {
166
			return record
167 168 169 170 171
		}
	}
	return nil
}

172 173
func (c *CapturingHandler) FindLogs(filters ...LogFilter) []*CapturedRecord {
	var logs []*CapturedRecord
174
	for _, record := range *c.Logs {
175 176 177 178 179 180 181 182
		match := true
		for _, filter := range filters {
			if !filter(record) {
				match = false
				break
			}
		}
		if match {
183
			logs = append(logs, record)
184 185 186 187 188
		}
	}
	return logs
}

189
func (h *CapturedRecord) AttrValue(name string) (v any) {
190 191 192 193
	h.Attrs(func(a slog.Attr) bool {
		if a.Key == name {
			v = a.Value.Any()
			return false
194
		}
195 196 197
		return true // try next
	})
	return
198 199
}

200
var _ slog.Handler = (*CapturingHandler)(nil)