add test to the vault, update deploy script, refactor contract

web3-nomad · Apr 10, 2024 · 539b05c · 539b05c
1 parent e61c0b7
commit 539b05c
Show file tree

Hide file tree

Showing 4 changed files with 187 additions and 100 deletions.
diff --git a/contracts/erc4626/Vault.sol b/contracts/erc4626/Vault.sol
@@ -16,18 +16,21 @@ contract HederaVault is IERC4626 {
     using Bits for uint256;
 
     ERC20 public immutable asset;
-    address newTokenAddress;
+    address public newTokenAddress;
     uint public totalTokens;
-    address[] tokenAddress;
+    address[] public tokenAddress;
     address public owner;
 
     event createdToken(address indexed createdTokenAddress);
 
     constructor(
         ERC20 _underlying,
         string memory _name,
-        string memory _symbol
+        string memory _symbol,
+        address[] memory rewardTokens
     ) payable ERC20(_name, _symbol, _underlying.decimals()) {
+        owner = msg.sender;
+
         SafeHTS.safeAssociateToken(address(_underlying), address(this));
         uint256 supplyKeyType;
         uint256 adminKeyType;
@@ -58,6 +61,8 @@ contract HederaVault is IERC4626 {
         newTokenAddress = SafeHTS.safeCreateFungibleToken(newToken, 0, _underlying.decimals());
         emit createdToken(newTokenAddress);
         asset = _underlying;
+
+        tokenAddress = rewardTokens;
     }
 
     struct UserInfo {
@@ -81,8 +86,6 @@ contract HederaVault is IERC4626 {
     function deposit(uint256 amount, address to) public override returns (uint256 shares) {
         require((shares = previewDeposit(amount)) != 0, "ZERO_SHARES");
 
-        asset.approve(address(this), amount);
-
         asset.safeTransferFrom(msg.sender, address(this), amount);
 
         totalTokens += amount;
@@ -223,9 +226,11 @@ contract HederaVault is IERC4626 {
                         REWARDS LOGIC
     //////////////////////////////////////////////////////////////*/
 
-    function addReward(address _token, uint _amount) internal {
+    function addReward(address _token, uint _amount) public {
         require(_amount != 0, "please provide amount");
         require(totalTokens != 0, "no token staked yet");
+        require(msg.sender == owner, "Only owner");
+
         uint perShareRewards;
         perShareRewards = _amount.mulDivDown(1, totalTokens);
         if (!rewardsAddress[_token].exist) {

diff --git a/scripts/deployVault.ts b/scripts/deployVault.ts
@@ -14,35 +14,46 @@ async function main() {
   let client = Client.forTestnet();
 
   const operatorPrKey = PrivateKey.fromStringECDSA(process.env.PRIVATE_KEY || '');
-  const operatorAccountId = AccountId.fromString(process.env.OPERATOR_ID || '');
+  const operatorAccountId = AccountId.fromString(process.env.ACCOUNT_ID || '');
 
   client.setOperator(
     operatorAccountId,
     operatorPrKey
   );
 
-  const createERC4626 = await createFungibleToken(
+  const stakingToken = await createFungibleToken(
     "ERC4626 on Hedera",
     "HERC4626",
-    process.env.OPERATOR_ID,
+    process.env.ACCOUNT_ID,
     operatorPrKey.publicKey,
     client,
     operatorPrKey
   );
 
-  const stakingTokenAddress = "0x" + createERC4626!.toSolidityAddress();
+  const rewardToken = await createFungibleToken(
+    "Reward Token 1",
+    "RT1",
+    process.env.ACCOUNT_ID,
+    operatorPrKey.publicKey,
+    client,
+    operatorPrKey
+  );
+
+  const stakingTokenAddress = "0x" + stakingToken!.toSolidityAddress();
+  const rewardTokenAddress = "0x" + rewardToken!.toSolidityAddress();
 
   const HederaVault = await ethers.getContractFactory("HederaVault");
   const hederaVault = await HederaVault.deploy(
     stakingTokenAddress,
     "TST",
     "TST",
-    { from: deployer.address, value: ethers.parseUnits("10", 18) }
+    [rewardTokenAddress],
+    { from: deployer.address, gasLimit: 3000000, value: ethers.parseUnits("10", 18) }
   );
   console.log("Hash ", hederaVault.deploymentTransaction()?.hash);
   await hederaVault.waitForDeployment();
 
-  console.log(await hederaVault.getAddress());
+  console.log("Vault deployed with address: ", await hederaVault.getAddress());
 }
 
 main().catch((error) => {

diff --git a/scripts/utils.ts b/scripts/utils.ts
@@ -85,7 +85,7 @@ export async function addToken(
     client: Client
 ) {
     let contractFunctionParameters = new ContractFunctionParameters()
-        .addAddress(tokenId.toSolidityAddress())
+        .addAddress(tokenId)
         .addUint256(amount * 1e8);
 
     const notifyRewardTx = await new ContractExecuteTransaction()
@@ -153,5 +153,7 @@ module.exports = {
     getClient,
     createAccount,
     mintToken,
-    TokenTransfer
+    TokenTransfer,
+    TokenBalance,
+    addToken,
 }