test(net): add message handler test cases

framework/src/test/java/org/tron/core/net/messagehandler/FetchInvDataMsgHandlerTest.java

@@ -1,5 +1,65 @@
 package org.tron.core.net.messagehandler;
 
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import java.lang.reflect.Field;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.tron.common.utils.ReflectUtils;
+import org.tron.common.utils.Sha256Hash;
+import org.tron.core.capsule.BlockCapsule;
+import org.tron.core.config.Parameter;
+import org.tron.core.net.TronNetDelegate;
+import org.tron.core.net.message.adv.BlockMessage;
+import org.tron.core.net.message.adv.FetchInvDataMessage;
+import org.tron.core.net.peer.Item;
+import org.tron.core.net.peer.PeerConnection;
+import org.tron.core.net.service.adv.AdvService;
+import org.tron.protos.Protocol;
+
+
+
 public class FetchInvDataMsgHandlerTest {
 
+  @Test
+  public void testProcessMessage() throws Exception {
+    FetchInvDataMsgHandler fetchInvDataMsgHandler = new FetchInvDataMsgHandler();
+    PeerConnection peer = Mockito.mock(PeerConnection.class);
+    TronNetDelegate tronNetDelegate = Mockito.mock(TronNetDelegate.class);
+    AdvService advService = Mockito.mock(AdvService.class);
+
+    Field field = FetchInvDataMsgHandler.class.getDeclaredField("tronNetDelegate");
+    field.setAccessible(true);
+    field.set(fetchInvDataMsgHandler, tronNetDelegate);
+
+    Mockito.when(tronNetDelegate.allowPBFT()).thenReturn(false);
+
+    BlockCapsule.BlockId blockId = new BlockCapsule.BlockId();
+    List<Sha256Hash> blockIds = new LinkedList<>();
+    blockIds.add(blockId);
+
+    Cache<Item, Long> advInvSpread = CacheBuilder.newBuilder().maximumSize(20000)
+        .expireAfterWrite(1, TimeUnit.HOURS).recordStats().build();
+    Mockito.when(peer.getAdvInvSpread()).thenReturn(advInvSpread);
+    Mockito.when(peer.isNeedSyncFromUs()).thenReturn(true);
+    Mockito.when(peer.isSyncFinish()).thenReturn(false);
+    Mockito.when(peer.getBlockBothHave()).thenReturn(blockId);
+    Cache<Sha256Hash, Long> syncBlockIdCache = CacheBuilder.newBuilder()
+        .maximumSize(2 * Parameter.NetConstants.SYNC_FETCH_BATCH_NUM).recordStats().build();
+    Mockito.when(peer.getSyncBlockIdCache()).thenReturn(syncBlockIdCache);
+    Mockito.when(peer.getLastSyncBlockId()).thenReturn(blockId);
+    BlockCapsule blockCapsule = new BlockCapsule(1, Sha256Hash.ZERO_HASH,
+        System.currentTimeMillis(), Sha256Hash.ZERO_HASH.getByteString());
+    Mockito.when(advService.getMessage(new Item(blockId, Protocol.Inventory.InventoryType.BLOCK)))
+        .thenReturn(new BlockMessage(blockCapsule));
+    ReflectUtils.setFieldValue(fetchInvDataMsgHandler, "advService", advService);
+
+    fetchInvDataMsgHandler.processMessage(peer,
+        new FetchInvDataMessage(blockIds, Protocol.Inventory.InventoryType.BLOCK));
+    Assert.assertNotNull(syncBlockIdCache.getIfPresent(blockId));
+  }
 }

framework/src/test/java/org/tron/core/net/messagehandler/PbftDataSyncHandlerTest.java

@@ -0,0 +1,57 @@
+package org.tron.core.net.messagehandler;
+
+import com.alibaba.fastjson.JSON;
+import com.google.protobuf.ByteString;
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.Map;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.tron.common.utils.Sha256Hash;
+import org.tron.core.ChainBaseManager;
+import org.tron.core.capsule.BlockCapsule;
+import org.tron.core.capsule.PbftSignCapsule;
+import org.tron.core.db.PbftSignDataStore;
+import org.tron.core.net.message.pbft.PbftCommitMessage;
+import org.tron.core.store.DynamicPropertiesStore;
+import org.tron.protos.Protocol;
+
+public class PbftDataSyncHandlerTest {
+  @Test
+  public void testProcessMessage() throws Exception {
+    PbftDataSyncHandler pbftDataSyncHandler = new PbftDataSyncHandler();
+    BlockCapsule blockCapsule = new BlockCapsule(1, Sha256Hash.ZERO_HASH,
+        System.currentTimeMillis(), ByteString.EMPTY);
+    Protocol.PBFTMessage.Raw.Builder rawBuilder = Protocol.PBFTMessage.Raw.newBuilder();
+    rawBuilder.setViewN(blockCapsule.getNum())
+        .setEpoch(0)
+        .setDataType(Protocol.PBFTMessage.DataType.BLOCK)
+        .setMsgType(Protocol.PBFTMessage.MsgType.PREPREPARE)
+        .setData(blockCapsule.getBlockId().getByteString());
+    Protocol.PBFTMessage.Raw raw = rawBuilder.build();
+    PbftSignCapsule pbftSignCapsule = new PbftSignCapsule(raw.toByteString(), new ArrayList<>());
+    PbftCommitMessage pbftCommitMessage = new PbftCommitMessage(pbftSignCapsule);
+    pbftDataSyncHandler.processMessage(null, pbftCommitMessage);
+    Assert.assertEquals(Protocol.PBFTMessage.Raw.parseFrom(
+        pbftCommitMessage.getPBFTCommitResult().getData()).getViewN(), 1);
+
+    DynamicPropertiesStore dynamicPropertiesStore = Mockito.mock(DynamicPropertiesStore.class);
+    PbftSignDataStore pbftSignDataStore = Mockito.mock(PbftSignDataStore.class);
+    ChainBaseManager chainBaseManager = Mockito.mock(ChainBaseManager.class);
+    Mockito.when(chainBaseManager.getDynamicPropertiesStore()).thenReturn(dynamicPropertiesStore);
+    Mockito.when(dynamicPropertiesStore.allowPBFT()).thenReturn(true);
+    Mockito.when(dynamicPropertiesStore.getMaintenanceTimeInterval()).thenReturn(600L);
+    Mockito.when(chainBaseManager.getPbftSignDataStore()).thenReturn(pbftSignDataStore);
+
+    Field field = PbftDataSyncHandler.class.getDeclaredField("chainBaseManager");
+    field.setAccessible(true);
+    field.set(pbftDataSyncHandler, chainBaseManager);
+
+    pbftDataSyncHandler.processPBFTCommitData(blockCapsule);
+    Field field1 = PbftDataSyncHandler.class.getDeclaredField("pbftCommitMessageCache");
+    field1.setAccessible(true);
+    Map map = JSON.parseObject(JSON.toJSONString(field1.get(pbftDataSyncHandler)), Map.class);
+    Assert.assertFalse(map.containsKey(0));
+  }
+}

framework/src/test/java/org/tron/core/net/messagehandler/PbftMsgHandlerTest.java

@@ -0,0 +1,121 @@
+package org.tron.core.net.messagehandler;
+
+import static org.mockito.Mockito.mock;
+
+import com.google.protobuf.ByteString;
+import java.io.File;
+import java.lang.reflect.Field;
+import java.net.InetSocketAddress;
+import java.util.ArrayList;
+import java.util.Collections;
+import org.bouncycastle.util.encoders.Hex;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.tron.common.application.TronApplicationContext;
+import org.tron.common.crypto.SignInterface;
+import org.tron.common.crypto.SignUtils;
+import org.tron.common.utils.FileUtil;
+import org.tron.common.utils.PublicMethod;
+import org.tron.common.utils.ReflectUtils;
+import org.tron.common.utils.Sha256Hash;
+import org.tron.consensus.base.Param;
+import org.tron.consensus.pbft.message.PbftMessage;
+import org.tron.core.Constant;
+import org.tron.core.capsule.BlockCapsule;
+import org.tron.core.config.DefaultConfig;
+import org.tron.core.config.args.Args;
+import org.tron.core.consensus.PbftBaseImpl;
+import org.tron.core.exception.P2pException;
+import org.tron.core.net.TronNetService;
+import org.tron.core.net.message.MessageTypes;
+import org.tron.core.net.peer.PeerConnection;
+import org.tron.core.net.peer.PeerManager;
+import org.tron.p2p.P2pConfig;
+import org.tron.p2p.base.Parameter;
+import org.tron.p2p.connection.Channel;
+import org.tron.protos.Protocol;
+
+
+public class PbftMsgHandlerTest {
+  private static TronApplicationContext context;
+  private PeerConnection peer;
+  private static String dbPath = "output-pbft-message-handler-test";
+
+
+  @BeforeClass
+  public static void init() {
+    Args.setParam(new String[] {"--output-directory", dbPath, "--debug"},
+        Constant.TEST_CONF);
+    context = new TronApplicationContext(DefaultConfig.class);
+
+    TronNetService tronNetService = context.getBean(TronNetService.class);
+    Parameter.p2pConfig = new P2pConfig();
+    ReflectUtils.setFieldValue(tronNetService, "p2pConfig", Parameter.p2pConfig);
+  }
+
+  @AfterClass
+  public static void destroy() {
+    Args.clearParam();
+    context.destroy();
+    FileUtil.deleteDir(new File(dbPath));
+  }
+
+  @Before
+  public void clearPeers() {
+    try {
+      Field field = PeerManager.class.getDeclaredField("peers");
+      field.setAccessible(true);
+      field.set(PeerManager.class, Collections.synchronizedList(new ArrayList<>()));
+    } catch (NoSuchFieldException | IllegalAccessException e) {
+      //ignore
+    }
+  }
+
+  @Test
+  public void testPbft() throws Exception {
+    InetSocketAddress a1 = new InetSocketAddress("127.0.0.1", 10001);
+    Channel c1 = mock(Channel.class);
+    Mockito.when(c1.getInetSocketAddress()).thenReturn(a1);
+    Mockito.when(c1.getInetAddress()).thenReturn(a1.getAddress());
+    PeerManager.add(context, c1);
+    Assert.assertEquals(1, PeerManager.getPeers().size());
+    Assert.assertFalse(c1.isDisconnect());
+
+    peer = PeerManager.getPeers().get(0);
+    BlockCapsule blockCapsule = new BlockCapsule(1, Sha256Hash.ZERO_HASH,
+        System.currentTimeMillis(), ByteString.EMPTY);
+    PbftMessage pbftMessage = new PbftMessage();
+    Protocol.PBFTMessage.Raw.Builder rawBuilder = Protocol.PBFTMessage.Raw.newBuilder();
+    Protocol.PBFTMessage.Builder builder = Protocol.PBFTMessage.newBuilder();
+    rawBuilder.setViewN(blockCapsule.getNum())
+        .setEpoch(0)
+        .setDataType(Protocol.PBFTMessage.DataType.BLOCK)
+        .setMsgType(Protocol.PBFTMessage.MsgType.PREPREPARE)
+        .setData(blockCapsule.getBlockId().getByteString());
+    Protocol.PBFTMessage.Raw raw = rawBuilder.build();
+    builder.setRawData(raw);
+    SignInterface sign = SignUtils.fromPrivate(Hex.decode(PublicMethod.getRandomPrivateKey()),
+        true);
+    builder.setSignature(ByteString.copyFrom(sign.Base64toBytes(sign.signHash(
+        Sha256Hash.hash(true, raw.toByteArray())))));
+    Protocol.PBFTMessage message = builder.build();
+    pbftMessage.setType(MessageTypes.PBFT_MSG.asByte());
+    pbftMessage.setPbftMessage(message);
+    pbftMessage.setData(message.toByteArray());
+    pbftMessage.setSwitch(blockCapsule.isSwitch());
+    Param.getInstance().setPbftInterface(context.getBean(PbftBaseImpl.class));
+    peer.setNeedSyncFromPeer(false);
+    //Mockito.doNothing().when(pbftMessage).analyzeSignature();
+    try {
+      context.getBean(PbftMsgHandler.class).processMessage(peer, pbftMessage);
+    } catch (P2pException e) {
+      Assert.assertEquals(P2pException.TypeEnum.BAD_MESSAGE, e.getType());
+    }
+
+    Assert.assertEquals(1, PeerManager.getPeers().size());
+  }
+}

framework/src/test/java/org/tron/core/net/messagehandler/TransactionsMsgHandlerTest.java

@@ -0,0 +1,84 @@
+package org.tron.core.net.messagehandler;
+
+import com.google.protobuf.Any;
+import com.google.protobuf.ByteString;
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import org.joda.time.DateTime;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.tron.common.BaseTest;
+import org.tron.common.utils.ByteArray;
+import org.tron.core.Constant;
+import org.tron.core.config.args.Args;
+import org.tron.core.net.TronNetDelegate;
+import org.tron.core.net.message.adv.TransactionMessage;
+import org.tron.core.net.message.adv.TransactionsMessage;
+import org.tron.core.net.peer.Item;
+import org.tron.core.net.peer.PeerConnection;
+import org.tron.core.net.service.adv.AdvService;
+import org.tron.protos.Protocol;
+import org.tron.protos.contract.BalanceContract;
+
+public class TransactionsMsgHandlerTest extends BaseTest {
+  @BeforeClass
+  public static void init() {
+    Args.setParam(new String[]{"--output-directory", dbPath(), "--debug"},
+        Constant.TEST_CONF);
+
+  }
+
+  @Test
+  public void testProcessMessage() {
+    TransactionsMsgHandler transactionsMsgHandler = new TransactionsMsgHandler();
+    try {
+      Assert.assertFalse(transactionsMsgHandler.isBusy());
+
+      transactionsMsgHandler.init();
+
+      PeerConnection peer = Mockito.mock(PeerConnection.class);
+      TronNetDelegate tronNetDelegate = Mockito.mock(TronNetDelegate.class);
+      AdvService advService = Mockito.mock(AdvService.class);
+
+      Field field = TransactionsMsgHandler.class.getDeclaredField("tronNetDelegate");
+      field.setAccessible(true);
+      field.set(transactionsMsgHandler, tronNetDelegate);
+
+      BalanceContract.TransferContract transferContract = BalanceContract.TransferContract
+          .newBuilder()
+          .setAmount(10)
+          .setOwnerAddress(ByteString.copyFrom(ByteArray.fromHexString("121212a9cf")))
+          .setToAddress(ByteString.copyFrom(ByteArray.fromHexString("232323a9cf"))).build();
+
+      long transactionTimestamp = DateTime.now().minusDays(4).getMillis();
+      Protocol.Transaction trx = Protocol.Transaction.newBuilder().setRawData(
+          Protocol.Transaction.raw.newBuilder().setTimestamp(transactionTimestamp)
+          .setRefBlockNum(1)
+          .addContract(
+              Protocol.Transaction.Contract.newBuilder()
+                  .setType(Protocol.Transaction.Contract.ContractType.TransferContract)
+                  .setParameter(Any.pack(transferContract)).build()).build())
+          .build();
+      Map<Item, Long> advInvRequest = new ConcurrentHashMap<>();
+      Item item = new Item(new TransactionMessage(trx).getMessageId(),
+          Protocol.Inventory.InventoryType.TRX);
+      advInvRequest.put(item, 0L);
+      Mockito.when(peer.getAdvInvRequest()).thenReturn(advInvRequest);
+
+      List<Protocol.Transaction> transactionList = new ArrayList<>();
+      transactionList.add(trx);
+      transactionsMsgHandler.processMessage(peer, new TransactionsMessage(transactionList));
+      Assert.assertNull(advInvRequest.get(item));
+      //Thread.sleep(10);
+    } catch (Exception e) {
+      Assert.fail();
+    } finally {
+      transactionsMsgHandler.close();
+    }
+  }
+}