coder · coadler · May 25, 2022 · May 25, 2022 · May 25, 2022 · May 25, 2022
diff --git a/sloggers/slogtest/t.go b/sloggers/slogtest/t.go
@@ -9,6 +9,7 @@ import (
 	"context"
 	"log"
 	"os"
+	"sync"
 	"testing"
 
 	"cdr.dev/slog"
@@ -28,26 +29,50 @@ type Options struct {
 	// IgnoreErrors causes the test logger to not fatal the test
 	// on Fatal and not error the test on Error or Critical.
 	IgnoreErrors bool
+	// SkipCleanup skips adding a t.Cleanup call that prevents the logger from
+	// logging after a test has exited. This is necessary because race
+	// conditions exist when t.Log is called concurrently of a test exiting. Set
+	// to true if you don't need this behavior.
+	SkipCleanup bool
 }
 
 // Make creates a Logger that writes logs to tb in a human readable format.
 func Make(tb testing.TB, opts *Options) slog.Logger {
 	if opts == nil {
 		opts = &Options{}
 	}
-	return slog.Make(testSink{
+
+	sink := &testSink{
 		tb:   tb,
 		opts: opts,
-	})
+	}
+	if !opts.SkipCleanup {
+		tb.Cleanup(func() {
+			sink.mu.Lock()
+			defer sink.mu.Unlock()
+			sink.testDone = true
+		})
+	}
+
+	return slog.Make(sink)
 }
 
 type testSink struct {
-	tb     testing.TB
-	opts   *Options
-	stdlib bool
+	tb       testing.TB
+	opts     *Options
+	mu       sync.RWMutex
+	testDone bool
 }
 
-func (ts testSink) LogEntry(ctx context.Context, ent slog.SinkEntry) {
+func (ts *testSink) LogEntry(ctx context.Context, ent slog.SinkEntry) {
+	ts.mu.RLock()
+	defer ts.mu.RUnlock()
+
+	// Don't log after the test this sink was created in has finished.
+	if ts.testDone {
+		return
+	}
+
 	// The testing package logs to stdout and not stderr.
 	s := entryhuman.Fmt(os.Stdout, ent)
 
@@ -65,12 +90,12 @@ func (ts testSink) LogEntry(ctx context.Context, ent slog.SinkEntry) {
 	}
 }
 
-func (ts testSink) Sync() {}
+func (ts *testSink) Sync() {}
 
 var ctx = context.Background()
 
 func l(t testing.TB) slog.Logger {
-	return Make(t, nil)
+	return Make(t, &Options{SkipCleanup: true})
 }
 
 // Debug logs the given msg and fields to t via t.Log at the debug level.

diff --git a/sloggers/slogtest/t_test.go b/sloggers/slogtest/t_test.go
@@ -45,18 +45,48 @@ func TestIgnoreErrors(t *testing.T) {
 	l.Fatal(bg, "hello")
 }
 
+func TestCleanup(t *testing.T) {
+	t.Parallel()
+
+	tb := &fakeTB{}
+	l := slogtest.Make(tb, &slogtest.Options{})
+
+	for _, fn := range tb.cleanups {
+		fn()
+	}
+
+	// This shoud not log since the logger was cleaned up.
+	l.Info(bg, "hello")
+	assert.Equal(t, "no logs", 0, tb.logs)
+}
+
+func TestSkipCleanup(t *testing.T) {
+	t.Parallel()
+
+	tb := &fakeTB{}
+	slogtest.Make(tb, &slogtest.Options{
+		SkipCleanup: true,
+	})
+
+	assert.Len(t, "no cleanups", 0, tb.cleanups)
+}
+
 var bg = context.Background()
 
 type fakeTB struct {
 	testing.TB
 
-	errors int
-	fatals int
+	logs     int
+	errors   int
+	fatals   int
+	cleanups []func()
 }
 
 func (tb *fakeTB) Helper() {}
 
-func (tb *fakeTB) Log(v ...interface{}) {}
+func (tb *fakeTB) Log(v ...interface{}) {
+	tb.logs++
+}
 
 func (tb *fakeTB) Error(v ...interface{}) {
 	tb.errors++
@@ -66,3 +96,7 @@ func (tb *fakeTB) Fatal(v ...interface{}) {
 	tb.fatals++
 	panic("")
 }
+
+func (tb *fakeTB) Cleanup(fn func()) {
+	tb.cleanups = append(tb.cleanups, fn)
+}