feat(test): add test for Wallet.estimateEnergy

Author: kongchi-ken

framework/src/test/java/org/tron/core/WalletTest.java

@@ -30,6 +30,8 @@
 import com.google.protobuf.ByteString;
 import java.io.File;
 import java.util.Arrays;
+
+import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import org.joda.time.DateTime;
 import org.junit.AfterClass;
@@ -54,6 +56,9 @@
 import org.tron.core.capsule.AccountCapsule;
 import org.tron.core.capsule.AssetIssueCapsule;
 import org.tron.core.capsule.BlockCapsule;
+import org.tron.core.capsule.BytesCapsule;
+import org.tron.core.capsule.CodeCapsule;
+import org.tron.core.capsule.ContractCapsule;
 import org.tron.core.capsule.DelegatedResourceCapsule;
 import org.tron.core.capsule.ExchangeCapsule;
 import org.tron.core.capsule.ProposalCapsule;
@@ -69,6 +74,7 @@
 import org.tron.core.store.DynamicPropertiesStore;
 import org.tron.core.utils.ProposalUtil.ProposalType;
 import org.tron.core.utils.TransactionUtil;
+import org.tron.core.vm.program.Program;
 import org.tron.protos.Protocol;
 import org.tron.protos.Protocol.AccountType;
 import org.tron.protos.Protocol.Block;
@@ -84,6 +90,7 @@
 import org.tron.protos.contract.BalanceContract;
 import org.tron.protos.contract.BalanceContract.TransferContract;
 import org.tron.protos.contract.Common;
+import org.tron.protos.contract.SmartContractOuterClass;
 
 
 @Slf4j
@@ -950,5 +957,92 @@ public void testGetAssetIssueByName() {
     chainBaseManager.getAssetIssueV2Store().delete(assetCapsule.createDbV2Key());
     chainBaseManager.getDynamicPropertiesStore().saveAllowSameTokenName(0);
   }
+
+  @Test
+  @SneakyThrows
+  public void testEstimateEnergy() {
+    dbManager.getDynamicPropertiesStore().put("ALLOW_TVM_TRANSFER_TRC10".getBytes(),
+        new BytesCapsule(ByteArray.fromHexString("0x01")));
+    String contractAddress = "0x1A622D84ed49f01045f5f1a5AfcEb9c57e9cC3ca";
+
+    SmartContractOuterClass.SmartContract smartContract =
+        SmartContractOuterClass.SmartContract.newBuilder().build();
+    ContractCapsule capsule = new ContractCapsule(smartContract);
+    dbManager.getContractStore().put(ByteArray.fromHexString(contractAddress), capsule);
+
+    String codeString = "608060405234801561001057600080fd5b50d3801561001d57600080fd5b50d2801561002a57600080fd5b50600436106100495760003560e01c806385bb7d69146100555761004a565b5b610052610073565b50005b61005d610073565b60405161006a91906100b9565b60405180910390f35b600080600090505b60028110156100a657808261009091906100d4565b915060018161009f91906100d4565b905061007b565b5090565b6100b38161012a565b82525050565b60006020820190506100ce60008301846100aa565b92915050565b60006100df8261012a565b91506100ea8361012a565b9250827fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff0382111561011f5761011e610134565b5b828201905092915050565b6000819050919050565b7f4e487b7100000000000000000000000000000000000000000000000000000000600052601160045260246000fdfea26474726f6e58221220f3d0198323c67293b97323c101e294e6d2cac7fb29555292675277e11c275a4b64736f6c63430008060033";
+    CodeCapsule codeCapsule = new CodeCapsule(ByteArray.fromHexString(codeString));
+    dbManager.getCodeStore().put(ByteArray.fromHexString(contractAddress), codeCapsule);
+
+    SmartContractOuterClass.TriggerSmartContract contract =
+        SmartContractOuterClass.TriggerSmartContract.newBuilder()
+            .setOwnerAddress(ByteString.copyFrom(ByteArray.fromHexString(OWNER_ADDRESS)))
+            .setContractAddress(ByteString.copyFrom(
+                ByteArray.fromHexString(
+                    contractAddress)))
+            .build();
+    TransactionCapsule trxCap = wallet.createTransactionCapsule(contract, ContractType.TriggerSmartContract);
+
+    GrpcAPI.TransactionExtention.Builder trxExtBuilder = GrpcAPI.TransactionExtention.newBuilder();
+    GrpcAPI.Return.Builder retBuilder = GrpcAPI.Return.newBuilder();
+    GrpcAPI.EstimateEnergyMessage.Builder estimateBuilder
+        = GrpcAPI.EstimateEnergyMessage.newBuilder();
+
+    try {
+      wallet.estimateEnergy(
+          contract, trxCap, trxExtBuilder, retBuilder, estimateBuilder);
+      Assert.fail();
+    } catch (ContractValidateException exception) {
+      assertEquals("this node does not support estimate energy", exception.getMessage());
+    }
+
+    Args.getInstance().setEstimateEnergy(true);
+
+    wallet.estimateEnergy(
+        contract, trxCap, trxExtBuilder, retBuilder, estimateBuilder);
+    GrpcAPI.EstimateEnergyMessage message = estimateBuilder.build();
+    Assert.assertTrue(message.getEnergyRequired() > 0);
+  }
+
+  @Test
+  @SneakyThrows
+  public void testEstimateEnergyOutOfTime() {
+    dbManager.getDynamicPropertiesStore().put("ALLOW_TVM_TRANSFER_TRC10".getBytes(),
+        new BytesCapsule(ByteArray.fromHexString("0x01")));
+
+    String contractAddress = "0x1A622D84ed49f01045f5f1a5AfcEb9c57e9cC3ca";
+
+    SmartContractOuterClass.SmartContract smartContract =
+        SmartContractOuterClass.SmartContract.newBuilder().build();
+    ContractCapsule capsule = new ContractCapsule(smartContract);
+    dbManager.getContractStore().put(ByteArray.fromHexString(contractAddress), capsule);
+
+    String codeString = "608060405234801561001057600080fd5b50d3801561001d57600080fd5b50d2801561002a57600080fd5b50600436106100495760003560e01c806385bb7d69146100555761004a565b5b610052610073565b50005b61005d610073565b60405161006a91906100ae565b60405180910390f35b600080600090505b64e8d4a5100081101561009b57808261009491906100c9565b915061007b565b5090565b6100a88161011f565b82525050565b60006020820190506100c3600083018461009f565b92915050565b60006100d48261011f565b91506100df8361011f565b9250827fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff0382111561011457610113610129565b5b828201905092915050565b6000819050919050565b7f4e487b7100000000000000000000000000000000000000000000000000000000600052601160045260246000fdfea26474726f6e58221220a7e1a6e6d17684029015a0b593b63440f77e7eb8abd4297a3063e59f28086bf464736f6c63430008060033";
+    CodeCapsule codeCapsule = new CodeCapsule(ByteArray.fromHexString(codeString));
+    dbManager.getCodeStore().put(ByteArray.fromHexString(contractAddress), codeCapsule);
+
+    SmartContractOuterClass.TriggerSmartContract contract =
+        SmartContractOuterClass.TriggerSmartContract.newBuilder()
+            .setOwnerAddress(ByteString.copyFrom(ByteArray.fromHexString(OWNER_ADDRESS)))
+            .setContractAddress(ByteString.copyFrom(
+                ByteArray.fromHexString(
+                    contractAddress)))
+            .build();
+    TransactionCapsule trxCap = wallet.createTransactionCapsule(contract, ContractType.TriggerSmartContract);
+
+    GrpcAPI.TransactionExtention.Builder trxExtBuilder = GrpcAPI.TransactionExtention.newBuilder();
+    GrpcAPI.Return.Builder retBuilder = GrpcAPI.Return.newBuilder();
+    GrpcAPI.EstimateEnergyMessage.Builder estimateBuilder
+        = GrpcAPI.EstimateEnergyMessage.newBuilder();
+
+    Args.getInstance().setEstimateEnergy(true);
+
+    try {
+      wallet.estimateEnergy(
+          contract, trxCap, trxExtBuilder, retBuilder, estimateBuilder);
+      Assert.fail("EstimateEnergy should throw exception!");
+    } catch (Program.OutOfTimeException ignored) {
+    }
+  }
 }