fix mocks

Signed-off-by: garyschulte <garyschulte@gmail.com>
hyperledger · Oct 10, 2022 · 12c045c · 12c045c
1 parent 3e85631
commit 12c045c
Showing 1 changed file with 10 additions and 18 deletions.
diff --git a/.../core/src/test/java/org/hyperledger/besu/ethereum/bonsai/BonsaiWorldStateArchiveTest.java b/.../core/src/test/java/org/hyperledger/besu/ethereum/bonsai/BonsaiWorldStateArchiveTest.java
@@ -173,12 +173,10 @@ public void testGetMutableWithStorageConsistencyNotRollbackTheState() {
                     layeredWorldStatesByHash),
                 storageProvider,
                 blockchain));
-    var updater =
-        spy(
-            bonsaiWorldStateArchive.getUpdaterFromPersistedState(
-                any(BonsaiSnapshotWorldState.class)));
-    when(bonsaiWorldStateArchive.getUpdaterFromPersistedState(any(BonsaiSnapshotWorldState.class)))
-        .thenReturn(updater);
+
+    var worldState = (BonsaiPersistedWorldState) bonsaiWorldStateArchive.getMutable();
+    var updater = spy(bonsaiWorldStateArchive.getUpdaterFromPersistedState(worldState));
+    when(bonsaiWorldStateArchive.getUpdaterFromPersistedState(worldState)).thenReturn(updater);
 
     final BlockHeader blockHeader = blockBuilder.number(0).buildHeader();
 
@@ -221,12 +219,9 @@ public void testGetMutableWithStorageConsistencyToRollbackAndRollForwardTheState
                     layeredWorldStatesByHash),
                 storageProvider,
                 blockchain));
-    var updater =
-        spy(
-            bonsaiWorldStateArchive.getUpdaterFromPersistedState(
-                any(BonsaiSnapshotWorldState.class)));
-    when(bonsaiWorldStateArchive.getUpdaterFromPersistedState(any(BonsaiSnapshotWorldState.class)))
-        .thenReturn(updater);
+    var worldState = (BonsaiPersistedWorldState) bonsaiWorldStateArchive.getMutable();
+    var updater = spy(bonsaiWorldStateArchive.getUpdaterFromPersistedState(worldState));
+    when(bonsaiWorldStateArchive.getUpdaterFromPersistedState(worldState)).thenReturn(updater);
 
     // initial persisted state hash key
     when(blockchain.getBlockHeader(eq(Hash.ZERO))).thenReturn(Optional.of(blockHeaderChainA));
@@ -276,12 +271,9 @@ public void testGetMutableWithRollbackNotOverrideTrieLogLayer() {
                     layeredWorldStatesByHash),
                 storageProvider,
                 blockchain));
-    var updater =
-        spy(
-            bonsaiWorldStateArchive.getUpdaterFromPersistedState(
-                any(BonsaiSnapshotWorldState.class)));
-    when(bonsaiWorldStateArchive.getUpdaterFromPersistedState(any(BonsaiSnapshotWorldState.class)))
-        .thenReturn(updater);
+    var worldState = (BonsaiPersistedWorldState) bonsaiWorldStateArchive.getMutable();
+    var updater = spy(bonsaiWorldStateArchive.getUpdaterFromPersistedState(worldState));
+    when(bonsaiWorldStateArchive.getUpdaterFromPersistedState(worldState)).thenReturn(updater);
 
     // initial persisted state hash key
     when(blockchain.getBlockHeader(eq(Hash.ZERO))).thenReturn(Optional.of(blockHeaderChainA));