From 1a8603caf884535d3295f4779a7e23d281adac64 Mon Sep 17 00:00:00 2001
From: AionJayT <jay.tseng@aion.network>
Date: Thu, 19 Mar 2020 10:48:44 -0400
Subject: [PATCH] Add test for kernel internal shutdown

---
 modAionImpl/build.gradle                      |  1 +
 .../blockchain/StandaloneBlockchainTest.java  | 30 ++++++++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/modAionImpl/build.gradle b/modAionImpl/build.gradle
index 7cbe2202b0..48f9653052 100644
--- a/modAionImpl/build.gradle
+++ b/modAionImpl/build.gradle
@@ -73,6 +73,7 @@ dependencies {
     testCompile group: 'org.apache.commons', name: 'commons-lang3', version: '3.4'
     testCompile "org.mockito:mockito-core:2.23.0"
     testCompile 'com.google.truth:truth:0.42'
+    testCompile 'com.github.stefanbirkner:system-rules:1.19.0'
 }
 
 task preBuild(type: Exec) {
diff --git a/modAionImpl/test/org/aion/zero/impl/blockchain/StandaloneBlockchainTest.java b/modAionImpl/test/org/aion/zero/impl/blockchain/StandaloneBlockchainTest.java
index 3ea1ec425e..2c0bc453e8 100644
--- a/modAionImpl/test/org/aion/zero/impl/blockchain/StandaloneBlockchainTest.java
+++ b/modAionImpl/test/org/aion/zero/impl/blockchain/StandaloneBlockchainTest.java
@@ -4,12 +4,19 @@
 
 import java.math.BigInteger;
 import java.util.Collections;
+import java.util.concurrent.TimeUnit;
 import org.aion.crypto.ECKey;
+import org.aion.db.impl.SystemExitCodes;
+import org.aion.mcf.blockchain.Block;
 import org.aion.types.AionAddress;
-import org.aion.zero.impl.types.AionBlock;
+import org.aion.zero.impl.core.ImportResult;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.contrib.java.lang.system.ExpectedSystemExit;
 
 public class StandaloneBlockchainTest {
+    @Rule
+    public final ExpectedSystemExit exit = ExpectedSystemExit.none();
 
     @Test
     public void testStandaloneBlockchainGenerateAccounts() {
@@ -25,4 +32,25 @@ public void testStandaloneBlockchainGenerateAccounts() {
                     .isGreaterThan(BigInteger.ZERO);
         }
     }
+
+    @Test
+    public void testBlockChainShutdownhook() {
+        exit.expectSystemExitWithStatus(SystemExitCodes.NORMAL);
+
+        StandaloneBlockchain.Builder builder = new StandaloneBlockchain.Builder();
+        StandaloneBlockchain.Bundle bundle =
+                builder.withValidatorConfiguration("simple").withDefaultAccounts().build();
+
+        StandaloneBlockchain sb = bundle.bc;
+
+        AionBlockchainImpl.shutdownHook = 2;
+        Block block1 = sb.createBlock(sb.genesis, Collections.EMPTY_LIST, false, TimeUnit.MILLISECONDS.toSeconds(System.currentTimeMillis()));
+        ImportResult result = sb.tryToConnect(block1);
+        assertThat(result.isBest()).isTrue();
+
+        Block block2 = sb.createBlock(block1, Collections.EMPTY_LIST, false, TimeUnit.MILLISECONDS.toSeconds(System.currentTimeMillis()));
+
+        // expect the system will be shutdown when import the block2
+        sb.tryToConnect(new BlockWrapper(block2));
+    }
 }