flumetest: use t.Cleanup() to revert changes

This means the caller doesn’t need to remember to defer call the returned cleanup function.  The caller can just call flumetest.Start(), and ignore the return value.  This is how flume v2 will work.

Also added more tests.

flumetest/flumetest.go

@@ -24,6 +24,7 @@ import (
 	"os"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"testing"
 )
 
@@ -119,6 +120,7 @@ func (l *lockedBuf) String() string {
 type testingTB interface {
 	Failed() bool
 	Log(args ...interface{})
+	Cleanup(func())
 }
 
 func start(t testingTB) func() {
@@ -129,11 +131,21 @@ func start(t testingTB) func() {
 		// need to use a synchronized version of buf, since
 		// logs may be written on other goroutines than this one,
 		// and bytes.Buffer is not concurrent safe.
-		buf := lockedBuf{
+		buf := &lockedBuf{
 			buf: bytes.NewBuffer(nil),
 		}
-		revertOut := flume.SetOut(&buf)
+		revertOut := flume.SetOut(buf)
+
+		// since we're calling this function via t.Cleanup *and* returning
+		// the function so the caller can call it with defer, there is a good
+		// chance it will be called twice.  I can't use sync.Once here, because
+		// if recover() is called inside the Once func, it doesn't work.  recover()
+		// must be called directly in the deferred function
+		ran := atomic.Bool{}
 		revert = func() {
+			if !ran.CompareAndSwap(false, true) {
+				return
+			}
 			revertOut()
 			// make sure that if the test panics or fails, we dump the logs
 			recovered := recover()
@@ -144,8 +156,24 @@ func start(t testingTB) func() {
 				panic(recovered)
 			}
 		}
-
 	}
 
+	t.Cleanup(revert)
+	// Calling Cleanup() to revert these changes should be sufficient, but isn't due to
+	// this bug: https://github.com/golang/go/issues/49929
+	// Due to that issue, if the test panics:
+	// 1. t.Failed() returns false inside the cleanup function
+	// 2. the revert doesn't know the test failed
+	// 3. the revert function doesn't flush its captured logs as it should when a test fails
+	//
+	// So we do both: call the revert function via t.Cleanup, as well as return a function
+	// that the test can call via defer.  t.Cleanup ensures we as least return the state
+	// of the system, even if the test itself doesn't call the revert cleanup function,
+	// but returning the cleanup function as well means tests that *do* call it via defer
+	// will correctly handle test panics.
+	//
+	// Since that bug is expected to be fixed soon in go, flume v2 may only rely on
+	// on t.Cleanup(), to make the API simpler, and live with the limitation until go fixes
+	// it.
 	return revert
 }

flumetest/flumetest_test.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"github.com/gemalto/flume"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"strings"
 	"sync"
 	"testing"
 )
@@ -13,11 +15,19 @@ func init() {
 }
 
 type mockT struct {
-	failed  bool
-	lastLog string
+	failed   bool
+	logs     strings.Builder
+	cleanups []func()
 	sync.Mutex
 }
 
+func (m *mockT) Cleanup(f func()) {
+	m.Lock()
+	defer m.Unlock()
+
+	m.cleanups = append(m.cleanups, f)
+}
+
 func (m *mockT) Failed() bool {
 	m.Lock()
 	defer m.Unlock()
@@ -27,52 +37,172 @@ func (m *mockT) Failed() bool {
 func (m *mockT) Log(args ...interface{}) {
 	m.Lock()
 	defer m.Unlock()
-	m.lastLog = fmt.Sprint(args...)
+	_, _ = fmt.Fprint(&m.logs, args...)
 }
 
 func TestStart(t *testing.T) {
 	var log = flume.New("TestStart")
 
-	fakeTestRun := func(succeed bool) string {
-		m := mockT{
-			failed: !succeed,
-		}
-		finish := start(&m)
+	tests := []struct {
+		name     string
+		failTest bool
+		testFunc func(tb testingTB)
+		expect   string
+		skip     string
+	}{
+		{
+			name: "success",
+			testFunc: func(tb testingTB) {
+				defer start(tb)()
 
-		log.Info("Hi", "color", "red")
+				log.Info("Hi", "color", "red")
+			},
+			failTest: false,
+			expect:   "",
+		},
+		{
+			name: "failed",
+			testFunc: func(tb testingTB) {
+				defer start(tb)()
 
-		finish()
-		m.Lock()
-		defer m.Unlock()
-		return m.lastLog
-	}
-	assert.Empty(t, fakeTestRun(true), "should not have logged, because the test didn't fail")
-	assert.Contains(t, fakeTestRun(false), "color:red", "should have logged since test failed")
+				log.Info("Hi", "color", "red")
+			},
+			failTest: true,
+			expect:   "color:red",
+		},
+		{
+			name:     "panic",
+			failTest: false,
+			expect:   "color:red",
+			testFunc: func(tb testingTB) {
+				require.Panics(t, func() {
+					defer start(tb)()
 
-	// this test is meant to trigger the race detector if we're not synchronizing correctly on the message
-	// buffer
-	m := mockT{
-		failed: true,
+					log.Info("Hi", "color", "red")
+
+					panic("boom")
+				})
+			},
+		},
+		{
+			name:     "race",
+			failTest: false,
+			expect:   "",
+			testFunc: func(tb testingTB) {
+				cleanup := start(tb)
+
+				// when run with the race detector, this would cause a race
+				// unless the log buffer is synchronized
+				barrier, stop := make(chan struct{}, 1), make(chan struct{})
+				go func() {
+					barrier <- struct{}{}
+					for {
+						select {
+						case <-stop:
+							return
+						default:
+							log.Info("Hi", "color", "red")
+						}
+					}
+				}()
+				<-barrier
+				cleanup()
+				stop <- struct{}{}
+			},
+		},
+		{
+			name:     "verbose",
+			failTest: false,
+			expect:   "color:red",
+			testFunc: func(tb testingTB) {
+				oldVerbose := Verbose
+				Verbose = true
+				defer func() {
+					Verbose = oldVerbose
+				}()
+				start(tb)
+
+				log.Info("Hi", "color", "red")
+			},
+		},
+		{
+			name:     "cleanup_without_defer",
+			failTest: true,
+			expect:   "color:red",
+			testFunc: func(tb testingTB) {
+				start(tb)
+
+				log.Info("Hi", "color", "red")
+			},
+		},
+		{
+			skip:     "this will fail until this golang issue is resolved: https://github.com/golang/go/issues/49929",
+			name:     "cleanup_without_defer_panic",
+			failTest: false,
+			expect:   "color:red",
+			testFunc: func(tb testingTB) {
+				start(tb)
+
+			},
+		},
 	}
-	finish := start(&m)
-
-	var wg sync.WaitGroup
-	wg.Add(1)
-	stop := make(chan struct{})
-	go func() {
-		wg.Done()
-		for {
-			select {
-			case <-stop:
-				return
-			default:
-				log.Info("logging on a different goroutine, for race detector")
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if test.skip != "" {
+				t.Skip(test.skip)
+			}
+
+			m := mockT{
+				failed: test.failTest,
+			}
+
+			test.testFunc(&m)
+
+			// call any registered cleanup functions, as the testing package would
+			// at the end of the test
+			for _, cleanup := range m.cleanups {
+				cleanup()
 			}
-		}
 
-	}()
-	wg.Wait()
-	finish()
-	stop <- struct{}{}
+			if test.expect == "" {
+				assert.Empty(t, m.logs.String())
+			} else {
+				assert.Contains(t, m.logs.String(), test.expect)
+			}
+		})
+	}
+
+	// fakeTestRun := func(succeed bool) string {
+	// 	m := mockT{
+	// 		failed: !succeed,
+	// 	}
+	// 	finish := start(&m)
+	//
+	// 	log.Info("Hi", "color", "red")
+	//
+	// 	finish()
+	// 	m.Lock()
+	// 	defer m.Unlock()
+	// 	return m.lastLog
+	// }
+	// assert.Empty(t, fakeTestRun(true), "should not have logged, because the test didn't fail")
+	// assert.Contains(t, fakeTestRun(false), "color:red", "should have logged since test failed")
+
+	// this test is meant to trigger the race detector if we're not synchronizing correctly on the message
+	// buffer
+	// m := mockT{
+	// 	failed: true,
+	// }
+	// finish := start(&m)
+	//
+	// barrier := make(chan struct{}, 1)
+	// go func() {
+	// 	barrier <- struct{}{}
+	// 	log.Info("logging on a different goroutine, for race detector")
+	//
+	// }()
+	// <-barrier
+	// finish()
 
 }