From 82bd6a8cc6e936167ff0201f2f56910dd96af7aa Mon Sep 17 00:00:00 2001
From: rian <rian.hughes@physics.ox.ac.uk>
Date: Tue, 24 Sep 2024 17:37:47 +0300
Subject: [PATCH] constrain Expect.RevertBlock() in test

---
 plugin/plugin_test.go |  8 +++++++-
 sync/sync.go          | 36 +++++++++++++++++++++++++++---------
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/plugin/plugin_test.go b/plugin/plugin_test.go
index b03ff94c3f..1c1d5d993c 100644
--- a/plugin/plugin_test.go
+++ b/plugin/plugin_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/NethermindEth/juno/clients/feeder"
 	"github.com/NethermindEth/juno/db/pebble"
 	"github.com/NethermindEth/juno/mocks"
+	junoplugin "github.com/NethermindEth/juno/plugin"
 	adaptfeeder "github.com/NethermindEth/juno/starknetdata/feeder"
 	"github.com/NethermindEth/juno/sync"
 	"github.com/NethermindEth/juno/utils"
@@ -55,7 +56,12 @@ func TestPlugin(t *testing.T) {
 		require.Equal(t, utils.HexToFelt(t, "0x34e815552e42c5eb5233b99de2d3d7fd396e575df2719bf98e7ed2794494f86"), head.Hash)
 
 		// Reorg 2 blocks, then sync 3 blocks
-		plugin.EXPECT().RevertBlock(gomock.Any(), gomock.Any(), gomock.Any()).Times(2)
+		su1, block1, err := integGw.StateUpdateWithBlock(context.Background(), uint64(1))
+		require.NoError(t, err)
+		su0, block0, err := integGw.StateUpdateWithBlock(context.Background(), uint64(0))
+		require.NoError(t, err)
+		plugin.EXPECT().RevertBlock(&junoplugin.BlockAndStateUpdate{block1, su1}, &junoplugin.BlockAndStateUpdate{block0, su0}, gomock.Any())
+		plugin.EXPECT().RevertBlock(&junoplugin.BlockAndStateUpdate{block0, su0}, &junoplugin.BlockAndStateUpdate{nil, nil}, gomock.Any())
 		for i := range 3 {
 			su, block, err := mainGw.StateUpdateWithBlock(context.Background(), uint64(i))
 			require.NoError(t, err)
diff --git a/sync/sync.go b/sync/sync.go
index 07e561a4cc..85fadf5c4e 100644
--- a/sync/sync.go
+++ b/sync/sync.go
@@ -3,6 +3,7 @@ package sync
 import (
 	"context"
 	"errors"
+	"fmt"
 	"runtime"
 	"sync/atomic"
 	"time"
@@ -188,21 +189,38 @@ func (s *Synchronizer) fetchUnknownClasses(ctx context.Context, stateUpdate *cor
 	return newClasses, closer()
 }
 
-func (s *Synchronizer) handlePluginRevertBlock(block *core.Block, stateUpdate *core.StateUpdate, reverseStateDiff *core.StateDiff) {
+func (s *Synchronizer) handlePluginRevertBlock() {
 	if s.plugin == nil {
 		return
 	}
-	toBlock, err := s.blockchain.Head()
+	block, err := s.blockchain.Head()
 	if err != nil {
 		s.log.Warnw("Failed to retrieve the reverted blockchain head block for the plugin", "err", err)
 		return
 	}
-
-	toSU, err := s.blockchain.StateUpdateByNumber(toBlock.Number)
+	stateUpdate, err := s.blockchain.StateUpdateByNumber(block.Number)
 	if err != nil {
 		s.log.Warnw("Failed to retrieve the reverted blockchain head state-update for the plugin", "err", err)
 		return
 	}
+	reverseStateDiff, err := s.blockchain.GetReverseStateDiff()
+	if err != nil {
+		s.log.Warnw("Failed to retrieve reverse state diff", "head", block.Number, "hash", block.Hash.ShortString(), "err", err)
+	}
+	var toBlock *core.Block
+	var toSU *core.StateUpdate
+	if block.Number != 0 {
+		toBlock, err = s.blockchain.BlockByHash(block.ParentHash)
+		if err != nil {
+			s.log.Warnw("Failed to retrieve the parent block for the plugin", "err", err)
+			return
+		}
+		toSU, err = s.blockchain.StateUpdateByNumber(toBlock.Number)
+		if err != nil {
+			s.log.Warnw("Failed to retrieve the parent state-update for the plugin", "err", err)
+			return
+		}
+	}
 	err = (s.plugin).RevertBlock(
 		&junoplugin.BlockAndStateUpdate{Block: block, StateUpdate: stateUpdate},
 		&junoplugin.BlockAndStateUpdate{Block: toBlock, StateUpdate: toSU},
@@ -231,20 +249,20 @@ func (s *Synchronizer) verifierTask(ctx context.Context, block *core.Block, stat
 				return
 			}
 			storeTimer := time.Now()
+			fmt.Println("Store(blo", block.Number)
 			err = s.blockchain.Store(block, commitments, stateUpdate, newClasses)
+			fmt.Println(err)
 			if err != nil {
 				if errors.Is(err, blockchain.ErrParentDoesNotMatchHead) {
 					// revert the head and restart the sync process, hoping that the reorg is not deep
 					// if the reorg is deeper, we will end up here again and again until we fully revert reorged
 					// blocks
-					reverseStateDiff, err := s.blockchain.GetReverseStateDiff()
-					if err != nil {
-						s.log.Warnw("Failed to retrieve reverse state diff", "head", block.Number, "hash", block.Hash.ShortString(), "err", err)
-					}
+					// Todo: move s.revertHead(block) before this
 					if s.plugin != nil {
-						s.handlePluginRevertBlock(block, stateUpdate, reverseStateDiff)
+						s.handlePluginRevertBlock()
 					}
 					s.revertHead(block)
+
 				} else {
 					s.log.Warnw("Failed storing Block", "number", block.Number,
 						"hash", block.Hash.ShortString(), "err", err)