fix test case

onflow · Dec 22, 2023 · 17b28f4 · 17b28f4
1 parent 4fe1c15
commit 17b28f4
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 17 deletions.
diff --git a/engine/access/rpc/backend/backend_accounts_test.go b/engine/access/rpc/backend/backend_accounts_test.go
@@ -376,7 +376,7 @@ func (s *BackendAccountsSuite) testGetAccountAtLatestBlock(ctx context.Context,
 
 func (s *BackendAccountsSuite) testGetAccountAtBlockHeight(ctx context.Context, backend *backendAccounts, statusCode codes.Code) {
 	height := s.block.Header.Height
-	s.headers.On("ByHeight", height).Return(s.block.Header, nil).Once()
+	s.headers.On("BlockIDByHeight", height).Return(s.block.Header.ID(), nil).Once()
 
 	if statusCode == codes.OK {
 		actual, err := backend.GetAccountAtBlockHeight(ctx, s.account.Address, height)

diff --git a/engine/consensus/approvals/testutil.go b/engine/consensus/approvals/testutil.go
@@ -134,20 +134,22 @@ func (s *BaseAssignmentCollectorTestSuite) SetupTest() {
 			return realstorage.ErrNotFound
 		}
 	})
-	s.Headers.On("ByHeight", mock.Anything).Return(
-		func(height uint64) *flow.Header {
+	s.Headers.On("BlockIDByHeight", mock.Anything).Return(
+		func(height uint64) (flow.Identifier, error) {
 			if block, found := s.FinalizedAtHeight[height]; found {
-				return block
+				return block.ID(), nil
 			} else {
-				return nil
+				return flow.ZeroID, realstorage.ErrNotFound
 			}
 		},
-		func(height uint64) error {
-			_, found := s.FinalizedAtHeight[height]
-			if !found {
-				return realstorage.ErrNotFound
+	)
+	s.Headers.On("ByHeight", mock.Anything).Return(
+		func(height uint64) (*flow.Header, error) {
+			if block, found := s.FinalizedAtHeight[height]; found {
+				return block, nil
+			} else {
+				return nil, realstorage.ErrNotFound
 			}
-			return nil
 		},
 	)
 

diff --git a/engine/execution/ingestion/loader/unexecuted_loader_test.go b/engine/execution/ingestion/loader/unexecuted_loader_test.go
@@ -200,7 +200,7 @@ func TestLoadingUnexecutedBlocks(t *testing.T) {
 		loader := loader.NewUnexecutedLoader(log, ps, headers, es)
 
 		// block C is the only finalized block, index its header by its height
-		headers.EXPECT().ByHeight(blockC.Header.Height).Return(blockC.Header, nil)
+		headers.EXPECT().BlockIDByHeight(blockC.Header.Height).Return(blockC.Header.ID(), nil)
 
 		es.ExecuteBlock(t, blockA)
 		es.ExecuteBlock(t, blockB)
@@ -238,7 +238,7 @@ func TestLoadingUnexecutedBlocks(t *testing.T) {
 		loader := loader.NewUnexecutedLoader(log, ps, headers, es)
 
 		// block C is finalized, index its header by its height
-		headers.EXPECT().ByHeight(blockC.Header.Height).Return(blockC.Header, nil)
+		headers.EXPECT().BlockIDByHeight(blockC.Header.Height).Return(blockC.Header.ID(), nil)
 
 		es.ExecuteBlock(t, blockA)
 		es.ExecuteBlock(t, blockB)
@@ -278,7 +278,7 @@ func TestLoadingUnexecutedBlocks(t *testing.T) {
 		loader := loader.NewUnexecutedLoader(log, ps, headers, es)
 
 		// block A is finalized, index its header by its height
-		headers.EXPECT().ByHeight(blockA.Header.Height).Return(blockA.Header, nil)
+		headers.EXPECT().BlockIDByHeight(blockA.Header.Height).Return(blockA.Header.ID(), nil)
 
 		es.ExecuteBlock(t, blockA)
 		es.ExecuteBlock(t, blockB)
@@ -346,7 +346,7 @@ func TestLoadingUnexecutedBlocks(t *testing.T) {
 		loader := loader.NewUnexecutedLoader(log, ps, headers, es)
 
 		// block C is finalized, index its header by its height
-		headers.EXPECT().ByHeight(blockC.Header.Height).Return(blockC.Header, nil)
+		headers.EXPECT().BlockIDByHeight(blockC.Header.Height).Return(blockC.Header.ID(), nil)
 
 		es.ExecuteBlock(t, blockA)
 		es.ExecuteBlock(t, blockB)

diff --git a/engine/execution/ingestion/loader/unfinalized_loader_test.go b/engine/execution/ingestion/loader/unfinalized_loader_test.go
@@ -37,9 +37,9 @@ func TestLoadingUnfinalizedBlocks(t *testing.T) {
 	es := new(stateMock.FinalizedExecutionState)
 	es.On("GetHighestFinalizedExecuted").Return(genesis.Header.Height)
 	headers := new(storage.Headers)
-	headers.On("ByHeight", blockA.Header.Height).Return(blockA.Header, nil)
-	headers.On("ByHeight", blockB.Header.Height).Return(blockB.Header, nil)
-	headers.On("ByHeight", blockC.Header.Height).Return(blockC.Header, nil)
+	headers.On("BlockIDByHeight", blockA.Header.Height).Return(blockA.Header.ID(), nil)
+	headers.On("BlockIDByHeight", blockB.Header.Height).Return(blockB.Header.ID(), nil)
+	headers.On("BlockIDByHeight", blockC.Header.Height).Return(blockC.Header.ID(), nil)
 
 	loader := loader.NewUnfinalizedLoader(unittest.Logger(), ps, headers, es)