fix(coderd/provisionerdserver): fix test flake in TestHeartbeat

johnstcn · johnstcn · commit bb32c587b000 · 2024-01-24T18:25:51.000Z
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -72,6 +72,9 @@ type Options struct {
 	// The default function just calls UpdateProvisionerDaemonLastSeenAt.
 	// This is mainly used for testing.
 	HeartbeatFn func(context.Context) error
+
+	// HeartbeatDone is used for testing.
+	HeartbeatDone chan struct{}
 }
 
 type server struct {
@@ -183,6 +186,9 @@ func NewServer(
 	if options.HeartbeatInterval == 0 {
 		options.HeartbeatInterval = DefaultHeartbeatInterval
 	}
+	if options.HeartbeatDone == nil {
+		options.HeartbeatDone = make(chan struct{})
+	}
 
 	s := &server{
 		lifecycleCtx:                lifecycleCtx,
@@ -213,7 +219,7 @@ func NewServer(
 		s.heartbeatFn = s.defaultHeartbeat
 	}
 
-	go s.heartbeatLoop()
+	go s.heartbeatLoop(options.HeartbeatDone)
 	return s, nil
 }
 
@@ -227,17 +233,21 @@ func (s *server) timeNow() time.Time {
 }
 
 // heartbeatLoop runs heartbeatOnce at the interval specified by HeartbeatInterval
-// until the lifecycle context is canceled.
-func (s *server) heartbeatLoop() {
+// until the lifecycle context is canceled. Done is closed on exit.
+func (s *server) heartbeatLoop(hbDone chan<- struct{}) {
 	tick := time.NewTicker(time.Nanosecond)
-	defer tick.Stop()
+	defer func() {
+		close(hbDone)
+	}()
 	for {
 		select {
 		case <-s.lifecycleCtx.Done():
 			s.Logger.Debug(s.lifecycleCtx, "heartbeat loop canceled")
+			tick.Stop()
 			return
 		case <-tick.C:
 			if s.lifecycleCtx.Err() != nil {
+				tick.Stop()
 				return
 			}
 			start := s.timeNow()
diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -104,6 +104,7 @@ func TestHeartbeat(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
 	heartbeatChan := make(chan struct{})
+	heartbeatDone := make(chan struct{})
 	heartbeatFn := func(hbCtx context.Context) error {
 		t.Logf("heartbeat")
 		select {
@@ -117,6 +118,7 @@ func TestHeartbeat(t *testing.T) {
 	//nolint:dogsled // ｡:ﾟ૮ ˶ˆ ﻌ ˆ˶ ა ﾟ:｡
 	_, _, _, _ = setup(t, false, &overrides{
 		ctx:               ctx,
+		heartbeatDone:     heartbeatDone,
 		heartbeatFn:       heartbeatFn,
 		heartbeatInterval: testutil.IntervalFast,
 	})
@@ -125,17 +127,9 @@ func TestHeartbeat(t *testing.T) {
 	require.True(t, ok, "first heartbeat not received")
 	_, ok = <-heartbeatChan
 	require.True(t, ok, "second heartbeat not received")
+	// Cancel the context. This should cause heartbeatDone to be closed.
 	cancel()
-	// Close the channel to ensure we don't receive any more heartbeats.
-	// The test will fail if we do.
-	defer func() {
-		if r := recover(); r != nil {
-			t.Fatalf("heartbeat received after cancel: %v", r)
-		}
-	}()
-
-	close(heartbeatChan)
-	<-time.After(testutil.IntervalMedium)
+	<-heartbeatDone
 }
 
 func TestAcquireJob(t *testing.T) {
@@ -1727,6 +1721,7 @@ type overrides struct {
 	acquireJobLongPollDuration  time.Duration
 	heartbeatFn                 func(ctx context.Context) error
 	heartbeatInterval           time.Duration
+	heartbeatDone               chan struct{}
 }
 
 func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisionerDaemonServer, database.Store, pubsub.Pubsub, database.ProvisionerDaemon) {
@@ -1751,6 +1746,9 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 	if ov.heartbeatInterval == 0 {
 		ov.heartbeatInterval = testutil.IntervalMedium
 	}
+	if ov.heartbeatDone == nil {
+		ov.heartbeatDone = make(chan struct{})
+	}
 	if ov.deploymentValues != nil {
 		deploymentValues = ov.deploymentValues
 	}
@@ -1815,6 +1813,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 			AcquireJobLongPollDur: pollDur,
 			HeartbeatInterval:     ov.heartbeatInterval,
 			HeartbeatFn:           ov.heartbeatFn,
+			HeartbeatDone:         ov.heartbeatDone,
 		},
 	)
 	require.NoError(t, err)
