test(pubsublite): fix flaky TestPeriodicTask (#4069)

Two executions of the task func can run before the first call to `Stop`. So we stop the task immediately after receiving the first value, to reduce the likelihood of this occurring. Note that flakes may still occur.

Fixes: #4064

pubsublite/internal/wire/periodic_task_test.go

@@ -20,48 +20,47 @@ import (
 )
 
 func TestPeriodicTask(t *testing.T) {
+	const pollInterval = 10 * time.Millisecond
 	var callCount int32
 	values := make(chan int32)
 	task := func() {
 		values <- atomic.AddInt32(&callCount, 1)
 	}
-	ptask := newPeriodicTask(10*time.Millisecond, task)
+	ptask := newPeriodicTask(pollInterval, task)
+	defer ptask.Stop()
 
-	t.Run("Start1", func(t *testing.T) {
+	t.Run("Start", func(t *testing.T) {
 		ptask.Start()
 		ptask.Start() // Tests duplicate start
 
-		if got, want := <-values, int32(1); got != want {
+		got := <-values
+
+		// Attempt to immediately stop the task after the first run.
+		// Note: if this test is still flaky, pollInterval can be increased.
+		ptask.Stop()
+
+		if want := int32(1); got != want {
 			t.Errorf("got %d, want %d", got, want)
 		}
 	})
 
-	t.Run("Stop1", func(t *testing.T) {
-		ptask.Stop()
-		ptask.Stop() // Tests duplicate stop
+	t.Run("Stop", func(t *testing.T) {
+		ptask.Stop() // Tests duplicate stop (also called in Start above)
 
+		// Wait at least the poll interval to ensure the task did not run.
+		time.Sleep(2 * pollInterval)
 		select {
 		case got := <-values:
 			t.Errorf("got unexpected value %d", got)
 		default:
 		}
 	})
 
-	t.Run("Start2", func(t *testing.T) {
+	t.Run("Restart", func(t *testing.T) {
 		ptask.Start()
 
 		if got, want := <-values, int32(2); got != want {
 			t.Errorf("got %d, want %d", got, want)
 		}
 	})
-
-	t.Run("Stop2", func(t *testing.T) {
-		ptask.Stop()
-
-		select {
-		case got := <-values:
-			t.Errorf("got unexpected value %d", got)
-		default:
-		}
-	})
 }