Run heartbeat leak tests without MySQL or test order

metrics/repl.lag/lag_test.go

@@ -7,8 +7,8 @@ import (
 	"database/sql"
 	"errors"
 	"strings"
+	"sync/atomic"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/assert"
 
@@ -121,25 +121,42 @@ func TestCollectWithAllSources(t *testing.T) {
 	}
 }
 
-const (
-	readerFunc    = "heartbeat.(*BlipReader).run"
-	readerTimeout = 5 * time.Second // > heartbeat.ReadTimeout: reader exits after an in-flight query
-)
-
-// stopReader stops the heartbeat reader, if any, when the test ends. Tests
-// that count reader goroutines require that earlier tests don't leak them.
+// stopReader stops the heartbeat reader, if any, when the test ends.
 func stopReader(t *testing.T, cleanup func()) {
 	if cleanup != nil {
 		t.Cleanup(cleanup)
 	}
 }
 
-// requireNoReaders waits for readers stopped by earlier tests to exit, so the
-// test can count only the readers that it starts.
-func requireNoReaders(t *testing.T) {
+// spyReader is a heartbeat reader that records whether it was stopped.
+type spyReader struct {
+	heartbeat.Reader
+	stopped atomic.Bool
+}
+
+func (r *spyReader) Stop() {
+	r.stopped.Store(true)
+	r.Reader.Stop()
+}
+
+// spyOnReader makes c wrap its heartbeat reader in a spyReader.
+func spyOnReader(c *Lag) {
+	c.newReader = func(args heartbeat.BlipReaderArgs) heartbeat.Reader {
+		return &spyReader{Reader: newBlipReader(args)}
+	}
+}
+
+// requireReaderStopped requires that Prepare made a reader and stopped it.
+// Counting goroutines isn't enough: the reader goroutine might not have
+// started yet.
+func requireReaderStopped(t *testing.T, c *Lag) {
 	t.Helper()
-	if got := test.WaitGoroutines(readerFunc, 0, readerTimeout); got != 0 {
-		t.Fatalf("%d heartbeat readers running before test, expected 0 (leaked by an earlier test?)", got)
+	spy, ok := c.lagReader.(*spyReader)
+	if !ok {
+		t.Fatalf("Prepare didn't make a heartbeat reader: %T", c.lagReader)
+	}
+	if !spy.stopped.Load() {
+		t.Error("Prepare didn't stop the heartbeat reader")
 	}
 }
 
@@ -172,37 +189,47 @@ func TestPrepareBlipAtMultipleLevelsReturnsReaderCleanup(t *testing.T) {
 	// Plans usually collect repl.lag at more than one level because levels
 	// inherit domains from lower levels. Prepare starts one reader for all
 	// levels, so it must return the cleanup func that stops that reader.
-	requireNoReaders(t)
-
 	c := NewLag(unreachableDB(t))
+	// Stop the reader if the test fails before cleanup, even if Prepare
+	// returned a nil cleanup func. Stop can be called more than once.
+	t.Cleanup(func() {
+		if c.lagReader != nil {
+			c.lagReader.Stop()
+		}
+	})
+	readers := test.NewGoroutines()
+	readers.Label()
 	cleanup, err := c.Prepare(context.Background(), blipPlan("fast", "medium", "slow"))
 	if err != nil {
 		t.Fatal(err)
 	}
-	if got := test.WaitGoroutines(readerFunc, 1, readerTimeout); got != 1 {
+	if got := readers.Wait(test.HeartbeatReaderFunc, 1, test.HeartbeatReaderTimeout); got != 1 {
 		t.Fatalf("%d heartbeat readers running after Prepare, expected 1", got)
 	}
 	if cleanup == nil {
 		t.Fatal("Prepare returned nil cleanup func, so the heartbeat reader can't be stopped")
 	}
 	cleanup()
-	if got := test.WaitGoroutines(readerFunc, 0, readerTimeout); got != 0 {
+	if got := readers.Wait(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); got != 0 {
 		t.Errorf("%d heartbeat readers running after cleanup, expected 0", got)
 	}
 }
 
 func TestPrepareBlipErrorStopsReader(t *testing.T) {
 	// The reader starts before the repl check is verified, which fails because
 	// the database is unreachable. Prepare must stop the reader it started.
-	requireNoReaders(t)
-
 	c := NewLag(unreachableDB(t))
+	spyOnReader(c)
 	plan := blipPlan("kpi")
 	plan.Levels["kpi"].Collect[DOMAIN].Options[OPT_REPL_CHECK] = "read_only"
-	if _, err := c.Prepare(context.Background(), plan); err == nil {
+	readers := test.NewGoroutines()
+	readers.Label()
+	_, err := c.Prepare(context.Background(), plan)
+	if err == nil {
 		t.Fatal("Prepare returned nil error, expected repl check error")
 	}
-	if got := test.WaitGoroutines(readerFunc, 0, readerTimeout); got != 0 {
+	requireReaderStopped(t, c)
+	if got := readers.Wait(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); got != 0 {
 		t.Errorf("%d heartbeat readers running after Prepare error, expected 0", got)
 	}
 }
@@ -248,11 +275,16 @@ func TestPrepareBlipAtMultipleLevelsDropsNoHeartbeat(t *testing.T) {
 					level.Collect[DOMAIN].Options[OPT_REPORT_NO_HEARTBEAT] = report
 				}
 			}
+			readers := test.NewGoroutines()
+			readers.Label()
 			cleanup, err := c.Prepare(context.Background(), plan)
-			stopReader(t, cleanup)
+			stopReader(t, cleanup) // backstop if the test fails before it stops the reader
 			if err != nil {
 				t.Fatal(err)
 			}
+			if got := readers.Wait(test.HeartbeatReaderFunc, 1, test.HeartbeatReaderTimeout); got != 1 {
+				t.Fatalf("%d heartbeat readers running after Prepare, expected 1 for all levels", got)
+			}
 			for level := range plan.Levels {
 				metrics, err := c.Collect(context.Background(), level)
 				if err != nil {
@@ -266,6 +298,10 @@ func TestPrepareBlipAtMultipleLevelsDropsNoHeartbeat(t *testing.T) {
 					t.Errorf("%s: got %+v, expected no metrics (dropped)", level, metrics)
 				}
 			}
+			cleanup() // stopReader calls it again, which is a no-op
+			if got := readers.Wait(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); got != 0 {
+				t.Errorf("%d heartbeat readers running after cleanup, expected 0", got)
+			}
 		})
 	}
 }

metrics/repl.lag/main_test.go

@@ -0,0 +1,23 @@
+// Copyright 2026 Block, Inc.
+
+package repllag
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/cashapp/blip/v2/test"
+)
+
+// TestMain fails the run if a heartbeat reader is still running after all
+// tests. Tests count only the readers that they start (test.Goroutines), so
+// this catches a reader leaked by a test that doesn't count readers.
+func TestMain(m *testing.M) {
+	code := m.Run()
+	if n := test.WaitGoroutines(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); n != 0 {
+		fmt.Fprintf(os.Stderr, "FAIL: %d heartbeat readers running after all tests, expected 0: a test didn't stop the reader that it started\n", n)
+		code = 1
+	}
+	os.Exit(code)
+}

monitor/engine_cleanup_test.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"sync"
 	"testing"
-	"time"
 
 	"github.com/cashapp/blip/v2"
 	"github.com/cashapp/blip/v2/metrics"
@@ -16,20 +15,11 @@ import (
 	"github.com/cashapp/blip/v2/test/mock"
 )
 
-const (
-	heartbeatReaderFunc    = "heartbeat.(*BlipReader).run"
-	heartbeatReaderTimeout = 5 * time.Second // > heartbeat.ReadTimeout: reader exits after an in-flight query
-)
-
 func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) {
-	_, db, err := test.Connection(test.DefaultMySQLVersion)
-	if err != nil {
-		if test.Build {
-			t.Skip("MySQL test fixture not running")
-		}
-		t.Fatal(err)
-	}
-	defer db.Close()
+	// Engine.Prepare only pings the database, and the repl.lag blip writer
+	// doesn't query it, so MySQL isn't needed
+	db := mock.PingDB()
+	t.Cleanup(func() { db.Close() }) // after the engines stop: cleanups run last-in, first-out
 
 	// repl.lag at two levels, like a plan where the slow level inherits
 	// domains from the fast level
@@ -46,39 +36,35 @@ func TestEnginePrepareStopDoesNotLeakHeartbeatReader(t *testing.T) {
 		},
 	}
 
-	// Count only readers started by this test: wait for readers stopped by
-	// earlier tests to exit
-	if got := test.WaitGoroutines(heartbeatReaderFunc, 0, heartbeatReaderTimeout); got != 0 {
-		t.Fatalf("%d heartbeat readers running before test, expected 0 (leaked by an earlier test?)", got)
-	}
+	// Count only readers started by this test
+	readers := test.NewGoroutines()
+	readers.Label()
 	for i := 0; i < 3; i++ {
 		engine := NewEngine(blip.ConfigMonitor{MonitorId: plan.MonitorId}, db)
+		t.Cleanup(engine.Stop) // if the test fails before Stop; a second Stop is a no-op
 		if err := engine.Prepare(context.Background(), plan, func() {}, func() {}); err != nil {
 			t.Fatal(err)
 		}
-		if got := test.WaitGoroutines(heartbeatReaderFunc, 1, heartbeatReaderTimeout); got != 1 {
+		if got := readers.Wait(test.HeartbeatReaderFunc, 1, test.HeartbeatReaderTimeout); got != 1 {
 			t.Fatalf("cycle %d: %d heartbeat readers running after Prepare, expected 1", i, got)
 		}
 		engine.Stop()
-		if got := test.WaitGoroutines(heartbeatReaderFunc, 0, heartbeatReaderTimeout); got != 0 {
+		if got := readers.Wait(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); got != 0 {
 			t.Fatalf("cycle %d: %d heartbeat readers running after Stop, expected 0", i, got)
 		}
 	}
 }
 
 func TestEnginePrepareErrorCleansUpPreparedCollectors(t *testing.T) {
-	_, db, err := test.Connection(test.DefaultMySQLVersion)
-	if err != nil {
-		if test.Build {
-			t.Skip("MySQL test fixture not running")
-		}
-		t.Fatal(err)
-	}
+	// Engine.Prepare only pings the database, and the mock collectors don't
+	// query it, so MySQL isn't needed
+	db := mock.PingDB()
 	defer db.Close()
 
 	// Two collectors. Domains are prepared in random order, so whichever is
-	// prepared second fails. Both return a cleanup func, which Prepare must
-	// call because it discards both collectors.
+	// prepared second fails. Prepare discards both collectors, so it must call
+	// the cleanup func of the first. The second returns a cleanup func with its
+	// error, which Prepare calls as a backstop (see blip.Collector.Prepare).
 	domains := []string{"test.cleanup-a", "test.cleanup-b"}
 	mux := &sync.Mutex{}
 	prepared := 0

monitor/main_test.go

@@ -0,0 +1,23 @@
+// Copyright 2026 Block, Inc.
+
+package monitor
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/cashapp/blip/v2/test"
+)
+
+// TestMain fails the run if a heartbeat reader is still running after all
+// tests. Tests count only the readers that they start (test.Goroutines), so
+// this catches a reader leaked by a test that doesn't count readers.
+func TestMain(m *testing.M) {
+	code := m.Run()
+	if n := test.WaitGoroutines(test.HeartbeatReaderFunc, 0, test.HeartbeatReaderTimeout); n != 0 {
+		fmt.Fprintf(os.Stderr, "FAIL: %d heartbeat readers running after all tests, expected 0: a test didn't stop the reader that it started\n", n)
+		code = 1
+	}
+	os.Exit(code)
+}