fix more tests

Signed-off-by: garyschulte <garyschulte@gmail.com>
garyschulte · Jul 23, 2023 · b428ba5 · b428ba5
1 parent 1dc6218
commit b428ba5
Showing 1 changed file with 24 additions and 18 deletions.
diff --git a/besu/src/test/java/org/hyperledger/besu/cli/subcommands/storage/StorageSubCommandTest.java b/besu/src/test/java/org/hyperledger/besu/cli/subcommands/storage/StorageSubCommandTest.java
@@ -23,13 +23,18 @@
 import static org.hyperledger.besu.ethereum.core.VariablesStorageHelper.getSampleVariableValues;
 import static org.hyperledger.besu.ethereum.core.VariablesStorageHelper.populateBlockchainStorage;
 import static org.hyperledger.besu.ethereum.core.VariablesStorageHelper.populateVariablesStorage;
+import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.BLOCKCHAIN;
+import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.VARIABLES;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
 import org.hyperledger.besu.cli.CommandTestAbstract;
-import org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier;
 import org.hyperledger.besu.services.kvstore.InMemoryKeyValueStorage;
+import org.hyperledger.besu.services.kvstore.SegmentedInMemoryKeyValueStorage;
+import org.hyperledger.besu.services.kvstore.SegmentedKeyValueStorageAdapter;
+
+import java.util.List;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -58,13 +63,14 @@ public void storageRevertVariablesSubCommandExists() {
 
   @Test
   public void revertVariables() {
-    final var kvVariables = new InMemoryKeyValueStorage();
-    final var kvBlockchain = new InMemoryKeyValueStorage();
-    when(rocksDBStorageFactory.create(eq(KeyValueSegmentIdentifier.VARIABLES), any(), any()))
-        .thenReturn(kvVariables);
-    when(rocksDBStorageFactory.create(eq(KeyValueSegmentIdentifier.BLOCKCHAIN), any(), any()))
-        .thenReturn(kvBlockchain);
-
+    final var kvVariablesSeg = new SegmentedInMemoryKeyValueStorage();
+    final var kvVariables = new SegmentedKeyValueStorageAdapter(VARIABLES, kvVariablesSeg);
+    final var kvBlockchainSeg = new SegmentedInMemoryKeyValueStorage();
+    final var kvBlockchain = new SegmentedKeyValueStorageAdapter(BLOCKCHAIN, kvBlockchainSeg);
+    when(rocksDBStorageFactory.create(eq(List.of(VARIABLES)), any(), any()))
+        .thenReturn(kvVariablesSeg);
+    when(rocksDBStorageFactory.create(eq(List.of(BLOCKCHAIN)), any(), any()))
+        .thenReturn(kvBlockchainSeg);
     final var variableValues = getSampleVariableValues();
     assertNoVariablesInStorage(kvBlockchain);
     populateVariablesStorage(kvVariables, variableValues);
@@ -77,12 +83,14 @@ public void revertVariables() {
 
   @Test
   public void revertVariablesWhenSomeVariablesDoNotExist() {
-    final var kvVariables = new InMemoryKeyValueStorage();
-    final var kvBlockchain = new InMemoryKeyValueStorage();
-    when(rocksDBStorageFactory.create(eq(KeyValueSegmentIdentifier.VARIABLES), any(), any()))
-        .thenReturn(kvVariables);
-    when(rocksDBStorageFactory.create(eq(KeyValueSegmentIdentifier.BLOCKCHAIN), any(), any()))
-        .thenReturn(kvBlockchain);
+    final var kvVariablesSeg = new SegmentedInMemoryKeyValueStorage();
+    final var kvVariables = new SegmentedKeyValueStorageAdapter(VARIABLES, kvVariablesSeg);
+    final var kvBlockchainSeg = new SegmentedInMemoryKeyValueStorage();
+    final var kvBlockchain = new SegmentedKeyValueStorageAdapter(BLOCKCHAIN, kvBlockchainSeg);
+    when(rocksDBStorageFactory.create(eq(List.of(VARIABLES)), any(), any()))
+        .thenReturn(kvVariablesSeg);
+    when(rocksDBStorageFactory.create(eq(List.of(BLOCKCHAIN)), any(), any()))
+        .thenReturn(kvBlockchainSeg);
 
     final var variableValues = getSampleVariableValues();
     variableValues.remove(FINALIZED_BLOCK_HASH);
@@ -100,10 +108,8 @@ public void revertVariablesWhenSomeVariablesDoNotExist() {
   public void doesNothingWhenVariablesAlreadyReverted() {
     final var kvVariables = new InMemoryKeyValueStorage();
     final var kvBlockchain = new InMemoryKeyValueStorage();
-    when(rocksDBStorageFactory.create(eq(KeyValueSegmentIdentifier.VARIABLES), any(), any()))
-        .thenReturn(kvVariables);
-    when(rocksDBStorageFactory.create(eq(KeyValueSegmentIdentifier.BLOCKCHAIN), any(), any()))
-        .thenReturn(kvBlockchain);
+    when(rocksDBStorageFactory.create(eq(VARIABLES), any(), any())).thenReturn(kvVariables);
+    when(rocksDBStorageFactory.create(eq(BLOCKCHAIN), any(), any())).thenReturn(kvBlockchain);
 
     final var variableValues = getSampleVariableValues();
     assertNoVariablesInStorage(kvVariables);