fix concurrent access in tx streaming tests

Author: peterargue

engine/access/rpc/backend/transactions/stream/stream_backend_test.go

@@ -2,7 +2,6 @@ package stream
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"os"
 	"testing"
@@ -51,6 +50,7 @@ import (
 	storagemock "github.com/onflow/flow-go/storage/mock"
 	"github.com/onflow/flow-go/storage/operation/pebbleimpl"
 	"github.com/onflow/flow-go/storage/store"
+	"github.com/onflow/flow-go/utils/concurrentmap"
 	"github.com/onflow/flow-go/utils/unittest"
 	"github.com/onflow/flow-go/utils/unittest/mocks"
 )
@@ -95,7 +95,7 @@ type TransactionStreamSuite struct {
 	sealedBlock    *flow.Block
 	finalizedBlock *flow.Block
 
-	blockMap map[uint64]*flow.Block
+	blockMap *concurrentmap.Map[uint64, *flow.Block]
 
 	txStreamBackend *TransactionStream
 
@@ -115,7 +115,7 @@ func TestTransactionStatusSuite(t *testing.T) {
 
 // SetupTest initializes the test dependencies, configurations, and mock objects for TransactionStreamSuite tests.
 func (s *TransactionStreamSuite) SetupTest() {
-	s.log = zerolog.New(zerolog.NewConsoleWriter())
+	s.log = unittest.Logger()
 	s.state = protocol.NewState(s.T())
 	s.sealedSnapshot = protocol.NewSnapshot(s.T())
 	s.finalSnapshot = protocol.NewSnapshot(s.T())
@@ -207,10 +207,9 @@ func (s *TransactionStreamSuite) initializeBackend() {
 
 	s.sealedBlock = s.rootBlock
 	s.finalizedBlock = unittest.BlockWithParentFixture(s.sealedBlock.ToHeader())
-	s.blockMap = map[uint64]*flow.Block{
-		s.sealedBlock.Height:    s.sealedBlock,
-		s.finalizedBlock.Height: s.finalizedBlock,
-	}
+	s.blockMap = concurrentmap.New[uint64, *flow.Block]()
+	s.blockMap.Add(s.sealedBlock.Height, s.sealedBlock)
+	s.blockMap.Add(s.finalizedBlock.Height, s.finalizedBlock)
 
 	txStatusDeriver := txstatus.NewTxStatusDeriver(
 		s.state,
@@ -344,45 +343,53 @@ func (s *TransactionStreamSuite) initializeBackend() {
 	)
 }
 
+func blockByID(blockMap *concurrentmap.Map[uint64, *flow.Block]) func(flow.Identifier) (*flow.Block, error) {
+	return func(blockID flow.Identifier) (*flow.Block, error) {
+		var block *flow.Block
+		blockMap.ForEach(func(height uint64, b *flow.Block) error {
+			if b.ID() == blockID {
+				block = b
+			}
+			return nil
+		})
+		if block == nil {
+			return nil, storage.ErrNotFound
+		}
+		return block, nil
+	}
+}
+
+func blockByHeight(blockMap *concurrentmap.Map[uint64, *flow.Block]) func(uint64) (*flow.Block, error) {
+	return func(height uint64) (*flow.Block, error) {
+		if block, ok := blockMap.Get(height); ok {
+			return block, nil
+		}
+		return nil, storage.ErrNotFound
+	}
+}
+
 // initializeMainMockInstructions sets up the main mock behaviors for components used in TransactionStreamSuite tests.
 func (s *TransactionStreamSuite) initializeMainMockInstructions() {
 	s.transactions.On("Store", mock.Anything).Return(nil).Maybe()
 
-	s.blocks.On("ByHeight", mock.AnythingOfType("uint64")).Return(mocks.StorageMapGetter(s.blockMap)).Maybe()
-	s.blocks.On("ByID", mock.Anything).Return(
-		func(blockID flow.Identifier) *flow.Block {
-			for _, block := range s.blockMap {
-				if block.ID() == blockID {
-					return block
-				}
-			}
-			return nil
-		},
-		func(blockID flow.Identifier) error {
-			for _, block := range s.blockMap {
-				if block.ID() == blockID {
-					return nil
-				}
-			}
-			return errors.New("block not found")
-		},
-	).Maybe()
+	s.blocks.On("ByHeight", mock.AnythingOfType("uint64")).Return(blockByHeight(s.blockMap)).Maybe()
+	s.blocks.On("ByID", mock.Anything).Return(blockByID(s.blockMap)).Maybe()
 
 	s.state.On("Final").Return(s.finalSnapshot, nil).Maybe()
-	s.state.On("AtBlockID", mock.AnythingOfType("flow.Identifier")).Return(func(blockID flow.Identifier) protocolint.Snapshot {
-		s.tempSnapshot.On("Head").Unset()
-		s.tempSnapshot.On("Head").Return(func() *flow.Header {
-			for _, block := range s.blockMap {
-				if block.ID() == blockID {
-					return block.ToHeader()
-				}
-			}
-
-			return nil
-		}, nil)
-
-		return s.tempSnapshot
-	}, nil).Maybe()
+	s.state.On("AtBlockID", mock.AnythingOfType("flow.Identifier")).Return(
+		func(blockID flow.Identifier) protocolint.Snapshot {
+			s.tempSnapshot.On("Head").Unset()
+			s.tempSnapshot.On("Head").Return(
+				func() (*flow.Header, error) {
+					block, err := blockByID(s.blockMap)(blockID)
+					if err != nil {
+						return nil, err
+					}
+					return block.ToHeader(), nil
+				}, nil)
+
+			return s.tempSnapshot
+		}, nil).Maybe()
 
 	s.finalSnapshot.On("Head").Return(func() *flow.Header {
 		return s.finalizedBlock.ToHeader()
@@ -439,7 +446,7 @@ func (s *TransactionStreamSuite) addNewFinalizedBlock(parent *flow.Header, notif
 		option(s.finalizedBlock)
 	}
 
-	s.blockMap[s.finalizedBlock.Height] = s.finalizedBlock
+	s.blockMap.Add(s.finalizedBlock.Height, s.finalizedBlock)
 
 	if notify {
 		s.broadcaster.Publish()