Merge pull request tronprotocol#1085 from tronprotocol/switch_back_when_switch_fork_failed

Switch back when switch fork failed

Author: zhaohong

src/test/java/org/tron/core/db/ManagerTest.java

@@ -315,11 +315,19 @@ public void doNotSwitch()
     dbManager.pushBlock(blockCapsule0);
     dbManager.pushBlock(blockCapsule1);
     context.getBean(KhaosDatabase.class).removeBlk(dbManager.getBlockIdByNum(num));
+    Exception exception = null;
     try {
       dbManager.pushBlock(blockCapsule2);
     } catch (NonCommonBlockException e) {
       logger.info("do not switch fork");
       Assert.assertNotNull(dbManager.getBlockStore().get(blockCapsule0.getBlockId().getBytes()));
+      Assert.assertEquals(blockCapsule0.getBlockId(),
+          dbManager.getBlockStore().get(blockCapsule0.getBlockId().getBytes()).getBlockId());
+      exception = e;
+    }
+
+    if (exception == null) {
+      throw new IllegalStateException();
     }
   }
 
@@ -365,11 +373,13 @@ public void switchBack()
     dbManager.pushBlock(blockCapsule1);
     try {
       dbManager.pushBlock(blockCapsule2);
-    } catch (Exception e) {
+    } catch (ValidateScheduleException e) {
       logger.info("the fork chain has error block");
     }
 
     Assert.assertNotNull(dbManager.getBlockStore().get(blockCapsule0.getBlockId().getBytes()));
+    Assert.assertEquals(blockCapsule0.getBlockId(),
+        dbManager.getBlockStore().get(blockCapsule0.getBlockId().getBytes()).getBlockId());
   }
 
   private Map<ByteString, String> addTestWitnessAndAccount() {

src/test/java/org/tron/core/db/RevokingStoreTest.java

@@ -64,6 +64,39 @@ public synchronized void testUndo() throws RevokingStoreIllegalStateException {
     Assert.assertTrue(revokingDatabase.getStack().isEmpty());
     Assert.assertTrue(tronDatabase.getDbSource().allKeys().isEmpty());
     Assert.assertEquals(revokingDatabase.getActiveDialog(), 0);
+
+    dialog = DialogOptional.instance().setValue(revokingDatabase.buildDialog());
+    revokingDatabase.disable();
+    TestProtoCapsule testProtoCapsule = new TestProtoCapsule("del".getBytes());
+    tronDatabase.put(testProtoCapsule.getData(), testProtoCapsule);
+    revokingDatabase.enable();
+
+    try (Dialog tmpDialog = revokingDatabase.buildDialog()) {
+      tronDatabase.put(testProtoCapsule.getData(), new TestProtoCapsule("del2".getBytes()));
+      tmpDialog.merge();
+    }
+
+    try (Dialog tmpDialog = revokingDatabase.buildDialog()) {
+      tronDatabase.put(testProtoCapsule.getData(), new TestProtoCapsule("del22".getBytes()));
+      tmpDialog.merge();
+    }
+
+    try (Dialog tmpDialog = revokingDatabase.buildDialog()) {
+      tronDatabase.put(testProtoCapsule.getData(), new TestProtoCapsule("del222".getBytes()));
+      tmpDialog.merge();
+    }
+
+    try (Dialog tmpDialog = revokingDatabase.buildDialog()) {
+      tronDatabase.delete(testProtoCapsule.getData());
+      tmpDialog.merge();
+    }
+
+    dialog.reset();
+
+    logger.info("**********testProtoCapsule:" + String.valueOf(tronDatabase.get(testProtoCapsule.getData())));
+    Assert.assertArrayEquals("del".getBytes(), tronDatabase.get(testProtoCapsule.getData()).getData());
+    Assert.assertEquals(testProtoCapsule, tronDatabase.get(testProtoCapsule.getData()));
+
     tronDatabase.close();
   }