[PDE-381] [test] [fix] increase coverage, remove holders (#132)

smart-wallets/src/extensions/AssetVault.sol

@@ -313,9 +313,6 @@ contract AssetVault is WalletUtils, IAssetVault {
      * @param amount Amount of AssetTokens included in this yield allowance
      * @param expiration Timestamp at which the yield expires
      */
-
-
-
     function acceptYieldAllowance(IAssetToken assetToken, uint256 amount, uint256 expiration) external {
         AssetVaultStorage storage $ = _getAssetVaultStorage();
         address beneficiary = msg.sender;

smart-wallets/src/interfaces/IYieldDistributionToken.sol

@@ -28,12 +28,12 @@ interface IYieldDistributionToken is IERC20 {
 
     /// @notice Indicates a failure because a yield deposit is made in the same block as the last one
     error DepositSameBlock();
-    
+
     /**
-    * @notice Indicates a failure because the transfer of CurrencyToken failed
-    * @param user Address of the user whose transfer failed
-    * @param currencyTokenAmount Amount of CurrencyToken to be transferred
-    */
+     * @notice Indicates a failure because the transfer of CurrencyToken failed
+     * @param user Address of the user whose transfer failed
+     * @param currencyTokenAmount Amount of CurrencyToken to be transferred
+     */
     error TransferFailed(address user, uint256 currencyTokenAmount);
 
     /**

smart-wallets/src/mocks/MockAssetToken.sol

@@ -53,7 +53,7 @@ contract MockAssetToken is IAssetToken, ERC20Upgradeable, OwnableUpgradeable {
 
     function pendingYield(
         address user
-    ) external override view returns (uint256) {
+    ) external view override returns (uint256) {
         // Mock implementation for testing
     }
 

smart-wallets/src/mocks/MockInvalidAssetToken.sol

@@ -54,7 +54,7 @@ contract MockInvalidAssetToken is IAssetToken {
 
     function pendingYield(
         address user
-    ) external override view returns (uint256) {
+    ) external view override returns (uint256) {
         // Mock implementation for testing
     }
 

smart-wallets/src/token/AssetToken.sol

@@ -29,10 +29,6 @@ contract AssetToken is WalletUtils, YieldDistributionToken, IAssetToken {
         uint256 totalValue;
         /// @dev Mapping of whitelisted users
         mapping(address user => bool whitelisted) isWhitelisted;
-        /// @dev List of all users that have ever held AssetTokens
-        address[] holders;
-        /// @dev Mapping of all users that have ever held AssetTokens
-        mapping(address user => bool held) hasHeld;
     }
 
     // keccak256(abi.encode(uint256(keccak256("plume.storage.AssetToken")) - 1)) & ~bytes32(uint256(0xff))
@@ -120,6 +116,10 @@ contract AssetToken is WalletUtils, YieldDistributionToken, IAssetToken {
         uint256 totalValue_,
         bool isWhitelistEnabled_
     ) YieldDistributionToken(owner, name, symbol, currencyToken, decimals_, tokenURI_) {
+        if (address(currencyToken) == address(0)) {
+            revert InvalidAddress();
+        }
+
         AssetTokenStorage storage $ = _getAssetTokenStorage();
         $.totalValue = totalValue_;
         isWhitelistEnabled = isWhitelistEnabled_;
@@ -293,22 +293,6 @@ contract AssetToken is WalletUtils, YieldDistributionToken, IAssetToken {
         return _getAssetTokenStorage().isWhitelisted[user];
     }
 
-    /// @notice List of all users that have ever held AssetTokens
-    function getHolders() external view returns (address[] memory) {
-        return _getAssetTokenStorage().holders;
-    }
-
-    /**
-     * @notice Check if the user has ever held AssetTokens
-     * @param user Address of the user to check
-     * @return held Boolean indicating if the user has ever held AssetTokens
-     */
-    function hasBeenHolder(
-        address user
-    ) external view returns (bool held) {
-        return _getAssetTokenStorage().hasHeld[user];
-    }
-
     /// @notice Price of an AssetToken based on its total value and total supply
     function getPricePerToken() external view returns (uint256) {
         return _getAssetTokenStorage().totalValue / totalSupply();

smart-wallets/test/AssetToken.t.sol

@@ -42,12 +42,6 @@ contract AssetTokenTest is Test {
     address public user2;
     address public walletProxyAddress;
 
-    // Events for testing
-    event Deposited(address indexed user, uint256 currencyTokenAmount);
-
-    // small hack to be excluded from coverage report
-    // function test_() public { }
-
     function setUp() public {
         owner = ADMIN_ADDRESS;
         user1 = address(0x1);
@@ -116,7 +110,6 @@ contract AssetTokenTest is Test {
         assertEq(assetToken.name(), "Asset Token", "Name mismatch");
         assertEq(assetToken.symbol(), "AT", "Symbol mismatch");
         assertEq(assetToken.decimals(), 18, "Decimals mismatch");
-        //assertEq(assetToken.tokenURI_(), "http://example.com/token", "TokenURI mismatch");
         assertEq(assetToken.totalSupply(), 10_000 * 10 ** 18, "Total supply mismatch");
         assertEq(assetToken.getTotalValue(), 10_000 * 10 ** 18, "Total value mismatch");
         assertFalse(assetToken.isWhitelistEnabled(), "Whitelist should be enabled");
@@ -184,18 +177,30 @@ contract AssetTokenTest is Test {
         vm.stopPrank();
     }
 
-    //TODO: convert to SmartWalletCall
     function test_GetBalanceAvailable() public {
+        // Create new MockSmartWallet instance
+        MockSmartWallet mockWallet = new MockSmartWallet();
+        uint256 totalBalance = 1000 * 10 ** 18;
+        uint256 lockedBalance = 300 * 10 ** 18;
+
         vm.startPrank(address(testWalletImplementation));
 
-        uint256 balance = 1000 * 10 ** 18;
-        assetToken.addToWhitelist(user1);
-        assetToken.mint(user1, balance);
+        // Setup the mock wallet with tokens
+        assetToken.addToWhitelist(address(mockWallet));
+        assetToken.mint(address(mockWallet), totalBalance);
 
-        assertEq(assetToken.getBalanceAvailable(user1), balance);
+        // Lock some tokens
+        mockWallet.lockTokens(IAssetToken(address(assetToken)), lockedBalance);
         vm.stopPrank();
-        // Note: To fully test getBalanceAvailable, you would need to mock a SmartWallet
-        // contract that implements the ISmartWallet interface and returns a locked balance.
+
+        // Test available balance
+        uint256 availableBalance = assetToken.getBalanceAvailable(address(mockWallet));
+
+        // Available balance should be total balance minus locked balance
+        assertEq(availableBalance, totalBalance - lockedBalance, "Available balance incorrect");
+        assertEq(
+            mockWallet.getBalanceLocked(IAssetToken(address(assetToken))), lockedBalance, "Locked balance incorrect"
+        );
     }
 
     function test_Transfer() public {
@@ -275,7 +280,6 @@ contract AssetTokenTest is Test {
     function isWhitelistEnabled() public view returns (bool) {
         return assetTokenWhitelisted.isWhitelistEnabled();
     }
-    // Update the test function
 
     function test_AddAndRemoveFromWhitelist() public {
         console.log("AssetToken owner:", assetTokenWhitelisted.owner());
@@ -319,28 +323,43 @@ contract AssetTokenTest is Test {
 
     function test_RequestYield() public {
         mockSmartWallet = new SmartWallet();
+        uint256 initialBalance = 100 ether;
+        uint256 yieldAmount = 10 ether;
 
+        // Setup: Mint tokens and deposit yield
         vm.startPrank(address(testWalletImplementation));
+
+        // Setup mock wallet
         assetToken.addToWhitelist(address(mockSmartWallet));
-        assetToken.mint(address(mockSmartWallet), 1000 * 10 ** 18);
+        assetToken.mint(address(mockSmartWallet), initialBalance);
+
+        // Setup yield
+        ERC20Mock(address(currencyToken)).mint(address(testWalletImplementation), yieldAmount);
+        currencyToken.approve(address(assetToken), yieldAmount);
+
+        // Deposit yield
+        vm.warp(block.timestamp + 1); // Advance time to avoid same block deposit
+        assetToken.depositYield(yieldAmount);
+
+        // Advance time to accrue yield
+        vm.warp(block.timestamp + 60); // Advance by 1 minute
         vm.stopPrank();
 
+        // Record balances before yield request
+        uint256 walletBalanceBefore = currencyToken.balanceOf(address(mockSmartWallet));
+
+        // Request yield
+        vm.expectEmit(true, true, true, false); // Don't check data as exact yield amount may vary
+        emit IYieldDistributionToken.YieldAccrued(address(mockSmartWallet), 0); // Amount will be checked separately
         assetToken.requestYield(address(mockSmartWallet));
-        // You may need to implement a way to verify that the yield was requested
-    }
 
-    function test_GetHoldersAndHasBeenHolder() public {
-        vm.startPrank(address(testWalletImplementation));
-        assetToken.addToWhitelist(user1);
-        assetToken.addToWhitelist(user2);
-        assetToken.mint(user1, 100 * 10 ** 18);
-        assetToken.mint(user2, 100 * 10 ** 18);
-        vm.stopPrank();
+        // Verify balances after yield request
+        uint256 walletBalanceAfter = currencyToken.balanceOf(address(mockSmartWallet));
+        uint256 claimedYield = walletBalanceAfter - walletBalanceBefore;
 
-        address[] memory holders = assetToken.getHolders();
-        assertEq(holders.length, 3, "Should have 3 holders (owner, user1, user2)");
-        assertTrue(assetToken.hasBeenHolder(user1), "User1 should be a holder");
-        assertTrue(assetToken.hasBeenHolder(user2), "User2 should be a holder");
+        // Assert balance changes
+        assertGt(claimedYield, 0, "Should have claimed some yield");
+        assertLe(claimedYield, yieldAmount, "Claimed yield should not exceed deposited amount");
     }
 
     function test_GetPricePerToken() public {
@@ -412,17 +431,14 @@ contract AssetTokenTest is Test {
 
         // Test: Successfully remove from whitelist
         vm.expectEmit(true, false, false, false);
-        emit AddressRemovedFromWhitelist(user1);
+        emit AssetToken.AddressRemovedFromWhitelist(user1);
         assetTokenWhitelisted.removeFromWhitelist(user1);
 
         // Verify user is no longer whitelisted
         assertFalse(assetTokenWhitelisted.isAddressWhitelisted(user1));
         vm.stopPrank();
     }
 
-    // Helper function to check if event was emitted
-    event AddressRemovedFromWhitelist(address indexed user);
-
     function test_DepositYield() public {
         uint256 depositAmount = 100 ether;
 
@@ -477,10 +493,9 @@ contract AssetTokenTest is Test {
         vm.warp(block.timestamp + 1);
         assetToken.depositYield(yieldAmount);
 
-        // Let yield accrue
-        vm.warp(block.timestamp + 1 days);
+        // Let yield accrue for just a few seconds
+        vm.warp(block.timestamp + 10); // Only 10 seconds instead of 1 hour
 
-        // Claim yield for both users
         vm.stopPrank();
 
         vm.prank(user1);
@@ -489,9 +504,16 @@ contract AssetTokenTest is Test {
         vm.prank(user2);
         (IERC20 token2, uint256 amount2) = assetToken.claimYield(user2);
 
+        // Debug logs
+        console.log("User 1 yield amount:", amount1);
+        console.log("User 2 yield amount:", amount2);
+        console.log("Total yield deposited:", yieldAmount);
+
         // Test assumptions:
         // 1. Both users should get roughly equal yield (within 1%)
         assertApproxEqRel(amount1, amount2, 0.01e18);
+        // 2. Total claimed yield should not exceed deposited amount
+        assertLe(amount1 + amount2, yieldAmount);
     }
 
     function test_YieldCalculationsWithMultipleDeposits() public {
@@ -510,24 +532,85 @@ contract AssetTokenTest is Test {
         vm.warp(block.timestamp + 1);
         assetToken.depositYield(firstYield);
 
-        // Second yield deposit
-        vm.warp(block.timestamp + 1 days);
+        // Second yield deposit after a short time
+        vm.warp(block.timestamp + 10); // Only 10 seconds
         ERC20Mock(address(currencyToken)).mint(address(testWalletImplementation), secondYield);
         currencyToken.approve(address(assetToken), secondYield);
         assetToken.depositYield(secondYield);
 
         vm.stopPrank();
 
-        // Claim yield
+        // Claim yield after a short time
         vm.startPrank(user1);
-        vm.warp(block.timestamp + 1 days);
+        vm.warp(block.timestamp + 10); // Only 10 seconds
 
         (IERC20 token, uint256 claimedAmount) = assetToken.claimYield(user1);
+
+        // Debug logs
+        console.log("Claimed amount:", claimedAmount);
+        console.log("Total yield deposited:", firstYield + secondYield);
+
         vm.stopPrank();
 
         // Test assumptions:
         // 1. Token should be the correct currency token
         assertEq(address(token), address(currencyToken));
+        // 2. Claimed amount should not exceed total deposited yield
+        assertLe(claimedAmount, firstYield + secondYield);
+    }
+
+    function test_RevertConstructorInvalidAddress() public {
+        // First test: Invalid owner (should revert with OwnableInvalidOwner)
+        vm.expectRevert(abi.encodeWithSignature("OwnableInvalidOwner(address)", address(0)));
+        new AssetToken(
+            address(0), // Invalid owner address
+            "Asset Token",
+            "AT",
+            ERC20(address(currencyToken)),
+            18,
+            "http://example.com/token",
+            1000 * 10 ** 18,
+            10_000 * 10 ** 18,
+            true
+        );
+
+        // Second test: Invalid currency token (should revert with InvalidAddress)
+        vm.expectRevert(AssetToken.InvalidAddress.selector);
+        new AssetToken(
+            address(this), // valid owner
+            "Asset Token",
+            "AT",
+            ERC20(address(0)), // Invalid currency token address
+            18,
+            "http://example.com/token",
+            1000 * 10 ** 18,
+            10_000 * 10 ** 18,
+            true
+        );
+    }
+
+    function test_GetBalanceAvailableWithLockedBalance() public {
+        uint256 initialBalance = 100 ether;
+        uint256 lockedBalance = 30 ether;
+
+        // Create new MockSmartWallet instance
+        MockSmartWallet mockWallet = new MockSmartWallet();
+
+        vm.startPrank(address(testWalletImplementation));
+
+        // Mint some tokens to the mock wallet
+        assetToken.mint(address(mockWallet), initialBalance);
+
+        // Lock some tokens
+        mockWallet.lockTokens(IAssetToken(address(assetToken)), lockedBalance);
+
+        vm.stopPrank();
+
+        // Check available balance
+        uint256 availableBalance = assetToken.getBalanceAvailable(address(mockWallet));
+
+        // Available balance should be initial balance minus locked balance
+        assertEq(availableBalance, initialBalance - lockedBalance);
     }
 
 }