Add tests to the Vault, deploy new Vault for negative cases, refactor contracts and tests

contracts/erc4626/IERC4626.sol

@@ -11,6 +11,8 @@ abstract contract IERC4626 is ERC20 {
     event Deposit(address indexed from, address indexed to, uint256 amount, uint256 shares);
     event Withdraw(address indexed from, address indexed to, uint256 amount, uint256 shares);
 
+    error ZeroShares(uint256 numberOfShares);
+
     /*///////////////////////////////////////////////////////////////
                             Mutable Functions
     //////////////////////////////////////////////////////////////*/

contracts/erc4626/Vault.sol

@@ -16,18 +16,26 @@ 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);
+    /**
+     * @notice CreatedToken event.
+     * @dev Emitted after contract initialization, when represented shares token is deployed.
+     *
+     * @param createdToken The address of created token.
+     */
+    event CreatedToken(address indexed createdToken);
 
     constructor(
         ERC20 _underlying,
         string memory _name,
         string memory _symbol
     ) payable ERC20(_name, _symbol, _underlying.decimals()) {
+        owner = msg.sender;
+
         SafeHTS.safeAssociateToken(address(_underlying), address(this));
         uint256 supplyKeyType;
         uint256 adminKeyType;
@@ -56,7 +64,7 @@ contract HederaVault is IERC4626 {
         newToken.expiry = expiry;
         newToken.tokenKeys = keys;
         newTokenAddress = SafeHTS.safeCreateFungibleToken(newToken, 0, _underlying.decimals());
-        emit createdToken(newTokenAddress);
+        emit CreatedToken(newTokenAddress);
         asset = _underlying;
     }
 
@@ -78,10 +86,15 @@ contract HederaVault is IERC4626 {
                         DEPOSIT/WITHDRAWAL LOGIC
     //////////////////////////////////////////////////////////////*/
 
+    /**
+     * @dev Deposits staking token to the Vault and returns shares.
+     *
+     * @param amount The amount of staking token to send.
+     * @param to The shares receiver address.
+     * @return shares The amount of shares to receive.
+     */
     function deposit(uint256 amount, address to) public override returns (uint256 shares) {
-        require((shares = previewDeposit(amount)) != 0, "ZERO_SHARES");
-
-        asset.approve(address(this), amount);
+        if ((shares = previewDeposit(amount)) == 0) revert ZeroShares(amount);
 
         asset.safeTransferFrom(msg.sender, address(this), amount);
 
@@ -96,6 +109,13 @@ contract HederaVault is IERC4626 {
         afterDeposit(amount);
     }
 
+    /**
+     * @dev Mints.
+     *
+     * @param shares The amount of shares to send.
+     * @param to The receiver of tokens.
+     * @return amount The amount of tokens to receive.
+     */
     function mint(uint256 shares, address to) public override returns (uint256 amount) {
         _mint(to, amount = previewMint(shares));
 
@@ -110,6 +130,14 @@ contract HederaVault is IERC4626 {
         afterDeposit(amount);
     }
 
+    /**
+     * @dev Withdraws staking token and burns shares.
+     *
+     * @param amount The amount of shares.
+     * @param to The staking token receiver.
+     * @param from The .
+     * @return shares The amount of shares to burn.
+     */
     function withdraw(uint256 amount, address to, address from) public override returns (uint256 shares) {
         beforeWithdraw(amount);
 
@@ -125,6 +153,14 @@ contract HederaVault is IERC4626 {
         asset.safeTransfer(to, amount);
     }
 
+    /**
+     * @dev Redeems .
+     *
+     * @param shares The amount of shares.
+     * @param to The staking token receiver.
+     * @param from The .
+     * @return amount The amount of shares to burn.
+     */
     function redeem(uint256 shares, address to, address from) public override returns (uint256 amount) {
         require((amount = previewRedeem(shares)) != 0, "ZERO_ASSETS");
 
@@ -141,17 +177,28 @@ contract HederaVault is IERC4626 {
                          INTERNAL HOOKS LOGIC
     //////////////////////////////////////////////////////////////*/
 
+    /**
+     * @dev Updates user state according to withdraw inputs.
+     *
+     * @param amount The amount of shares.
+     */
     function beforeWithdraw(uint256 amount) internal {
         // claimAllReward(0);
         userContribution[msg.sender].num_shares -= amount;
         totalTokens -= amount;
     }
 
+    /**
+     * @dev Updates user state according to withdraw inputs.
+     *
+     * @param amount The amount of shares.
+     */
     function afterDeposit(uint256 amount) internal {
         if (!userContribution[msg.sender].exist) {
             for (uint i; i < tokenAddress.length; i++) {
                 address token = tokenAddress[i];
                 userContribution[msg.sender].lastClaimedAmountT[token] = rewardsAddress[token].amount;
+                SafeHTS.safeAssociateToken(token, msg.sender);
             }
             userContribution[msg.sender].num_shares = amount;
             userContribution[msg.sender].exist = true;
@@ -223,15 +270,18 @@ contract HederaVault is IERC4626 {
                         REWARDS LOGIC
     //////////////////////////////////////////////////////////////*/
 
-    function addReward(address _token, uint _amount) internal {
+    function addReward(address _token, uint _amount) public payable {
         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) {
             tokenAddress.push(_token);
             rewardsAddress[_token].exist = true;
             rewardsAddress[_token].amount = perShareRewards;
+            SafeHTS.safeAssociateToken(_token, address(this));
             ERC20(_token).safeTransferFrom(address(msg.sender), address(this), _amount);
         } else {
             rewardsAddress[_token].amount += perShareRewards;

hardhat.config.ts

@@ -1,6 +1,7 @@
 import { HardhatUserConfig } from "hardhat/config";
 import "@nomicfoundation/hardhat-toolbox";
 import "@openzeppelin/hardhat-upgrades";
+import "@nomicfoundation/hardhat-chai-matchers";
 
 import * as dotenv from "dotenv";
 

scripts/deployVault.ts

@@ -14,35 +14,35 @@ 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 stakingTokenAddress = "0x" + stakingToken!.toSolidityAddress();
 
   const HederaVault = await ethers.getContractFactory("HederaVault");
   const hederaVault = await HederaVault.deploy(
     stakingTokenAddress,
     "TST",
     "TST",
-    { from: deployer.address, value: ethers.parseUnits("10", 18) }
+    { from: deployer.address, gasLimit: 3000000, value: ethers.parseUnits("12", 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) => {

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,
 }