fix flaky pipeline tests

Author: peterargue

module/executiondatasync/optimistic_sync/pipeline_functional_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/jordanschalm/lockctx"
@@ -436,9 +437,14 @@ func (p *PipelineFunctionalSuite) TestMainCtxCancellationDuringWaitingPersist()
 
 // TestPipelineShutdownOnParentAbandon verifies that the pipeline transitions correctly to a shutdown state when the parent is abandoned.
 func (p *PipelineFunctionalSuite) TestPipelineShutdownOnParentAbandon() {
+	assertNoError := func(err error) {
+		p.Require().NoError(err)
+	}
+
 	tests := []struct {
 		name        string
 		config      PipelineConfig
+		checkError  func(err error)
 		customSetup func(pipeline Pipeline, updateChan chan State, errChan chan error)
 	}{
 		{
@@ -449,43 +455,68 @@ func (p *PipelineFunctionalSuite) TestPipelineShutdownOnParentAbandon() {
 				},
 				parentState: StateAbandoned,
 			},
+			checkError:  assertNoError,
+			customSetup: func(pipeline Pipeline, updateChan chan State, errChan chan error) {},
 		},
 		{
 			name: "from StateProcessing",
+			config: PipelineConfig{
+				beforePipelineRun: func(pipeline *PipelineImpl) {
+					p.execDataRequester.On("RequestExecutionData", mock.Anything).Return(func(ctx context.Context) (*execution_data.BlockExecutionData, error) {
+						pipeline.OnParentStateUpdated(StateAbandoned) // abandon during processing step
+						return p.expectedExecutionData, nil
+					}).Once()
+					// this method may not be called depending on how quickly the RequestExecutionData
+					// mock returns.
+					p.txResultErrMsgsRequester.On("Request", mock.Anything).Return(p.expectedTxResultErrMsgs, nil).Maybe()
+				},
+				parentState: StateWaitingPersist,
+			},
+			checkError: func(err error) {
+				// depending on the timing, the error may be during or after the indexing step.
+				if err != nil {
+					p.Require().ErrorContains(err, "could not perform indexing")
+				} else {
+					p.Require().NoError(err)
+				}
+			},
 			customSetup: func(pipeline Pipeline, updateChan chan State, errChan chan error) {
-				waitForStateUpdates(p.T(), updateChan, errChan, StateProcessing)
-
-				pipeline.OnParentStateUpdated(StateAbandoned)
+				synctestWaitForStateUpdates(p.T(), updateChan, StateProcessing)
 			},
-			config: p.config,
 		},
 		{
 			name: "from StateWaitingPersist",
+			config: PipelineConfig{
+				beforePipelineRun: func(pipeline *PipelineImpl) {
+					p.execDataRequester.On("RequestExecutionData", mock.Anything).Return(p.expectedExecutionData, nil).Once()
+					p.txResultErrMsgsRequester.On("Request", mock.Anything).Return(p.expectedTxResultErrMsgs, nil).Once()
+				},
+				parentState: StateWaitingPersist,
+			},
+			checkError: assertNoError,
 			customSetup: func(pipeline Pipeline, updateChan chan State, errChan chan error) {
-				waitForStateUpdates(p.T(), updateChan, errChan, StateProcessing, StateWaitingPersist)
-
+				synctestWaitForStateUpdates(p.T(), updateChan, StateProcessing, StateWaitingPersist)
 				pipeline.OnParentStateUpdated(StateAbandoned)
 			},
-			config: p.config,
 		},
 	}
 
 	for _, test := range tests {
 		p.T().Run(test.name, func(t *testing.T) {
-			p.WithRunningPipeline(func(pipeline Pipeline, updateChan chan State, errChan chan error, cancel context.CancelFunc) {
-				p.execDataRequester.On("RequestExecutionData", mock.Anything).Return(p.expectedExecutionData, nil).Maybe()
-				p.txResultErrMsgsRequester.On("Request", mock.Anything).Return(p.expectedTxResultErrMsgs, nil).Maybe()
+			p.execDataRequester.On("RequestExecutionData", mock.Anything).Unset()
+			p.txResultErrMsgsRequester.On("Request", mock.Anything).Unset()
 
-				if test.customSetup != nil {
+			synctest.Test(p.T(), func(t *testing.T) {
+				p.WithRunningPipeline(func(pipeline Pipeline, updateChan chan State, errChan chan error, cancel context.CancelFunc) {
 					test.customSetup(pipeline, updateChan, errChan)
-				}
 
-				waitForStateUpdates(p.T(), updateChan, errChan, StateAbandoned)
-				waitForError(p.T(), errChan, nil)
+					synctestWaitForStateUpdates(p.T(), updateChan, StateAbandoned)
+					test.checkError(<-errChan)
 
-				p.Assert().Equal(StateAbandoned, pipeline.GetState())
-				p.Assert().Nil(p.core.workingData)
-			}, test.config)
+					p.Assert().Equal(StateAbandoned, pipeline.GetState())
+					p.Assert().Nil(p.core.workingData)
+				}, test.config)
+			})
 		})
 	}
 }
@@ -534,13 +565,18 @@ func (p *PipelineFunctionalSuite) WithRunningPipeline(
 	pipelineIsReady := make(chan struct{})
 
 	go func() {
+		defer close(errChan)
+
 		if pipelineConfig.beforePipelineRun != nil {
 			pipelineConfig.beforePipelineRun(pipeline)
 		}
 
 		close(pipelineIsReady)
 
-		errChan <- pipeline.Run(ctx, p.core, pipelineConfig.parentState)
+		err := pipeline.Run(ctx, p.core, pipelineConfig.parentState)
+		if err != nil {
+			errChan <- err
+		}
 	}()
 
 	<-pipelineIsReady

module/executiondatasync/optimistic_sync/pipeline_test_utils.go

@@ -2,6 +2,7 @@ package optimistic_sync
 
 import (
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/rs/zerolog"
@@ -116,3 +117,28 @@ func createPipeline(t *testing.T) (*PipelineImpl, *osmock.Core, <-chan State, *m
 
 	return pipeline, mockCore, stateReceiver.updateChan, parent
 }
+
+// synctestWaitForStateUpdates waits for a sequence of state updates to occur using synctest.Wait.
+// updates must be received in the correct order or the test will fail.
+// TODO: refactor all tests to use the synctest approach.
+func synctestWaitForStateUpdates(t *testing.T, updateChan <-chan State, expectedStates ...State) {
+	for _, expected := range expectedStates {
+		synctest.Wait()
+		update, ok := <-updateChan
+		require.True(t, ok, "update channel closed unexpectedly")
+		assert.Equalf(t, expected, update, "expected pipeline to transition to %s, but got %s", expected, update)
+	}
+}
+
+// synctestWaitForError waits for an error from the errChan using synctest.Wait and asserts it matches the expected error.
+// TODO: refactor all tests to use the synctest approach.
+func synctestWaitForError(t *testing.T, errChan <-chan error, expectedErr error) {
+	synctest.Wait()
+	err := <-errChan
+	if expectedErr == nil {
+		assert.NoError(t, err, "Pipeline should complete without errors")
+	} else {
+		assert.ErrorIs(t, err, expectedErr)
+	}
+
+}

module/state_synchronization/indexer/indexer_core.go

@@ -222,12 +222,14 @@ func (c *IndexerCore) IndexBlockData(data *execution_data.BlockExecutionDataEnti
 	g.Go(func() error {
 		start := time.Now()
 
-		// index all collections except the system chunk
 		// Note: the access ingestion engine also indexes collections, starting when the block is
 		// finalized. This process can fall behind due to the node being offline, resource issues
 		// or network congestion. This indexer ensures that collections are never farther behind
 		// than the latest indexed block. Calling the collection handler with a collection that
 		// has already been indexed is a noop.
+
+		// index all collections except the system chunk. if there is only a single chunk, it is the
+		// system chunk and can be skipped.
 		indexedCount := 0
 		if len(data.ChunkExecutionDatas) > 1 {
 			for _, chunk := range data.ChunkExecutionDatas[0 : len(data.ChunkExecutionDatas)-1] {