fix: total assets (#17)

* fix: calculate totalAssets * test: add test to ensure doesn't revert on zero
bgd-labs · Sep 6, 2024 · 687022d · 687022d
1 parent 111ba64
commit 687022d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/src/periphery/contracts/static-a-token/ERC4626StataTokenUpgradeable.sol b/src/periphery/contracts/static-a-token/ERC4626StataTokenUpgradeable.sol
@@ -132,6 +132,11 @@ abstract contract ERC4626StataTokenUpgradeable is ERC4626Upgradeable, IERC4626St
     return convertToAssets(maxRedeem(owner));
   }
 
+  ///@inheritdoc IERC4626
+  function totalAssets() public view override returns (uint256) {
+    return _convertToAssets(totalSupply(), Math.Rounding.Floor);
+  }
+
   ///@inheritdoc IERC4626
   function maxRedeem(address owner) public view override returns (uint256) {
     DataTypes.ReserveData memory reserveData = POOL.getReserveDataExtended(asset());

diff --git a/tests/periphery/static-a-token/ERC4626StataTokenUpgradeable.t.sol b/tests/periphery/static-a-token/ERC4626StataTokenUpgradeable.t.sol
@@ -60,6 +60,10 @@ contract ERC4626StataTokenUpgradeableTest is TestnetProcedures {
     assertEq(erc4626Upgradeable.previewRedeem(shares), assets);
   }
 
+  function test_totalAssets_shouldbeZeroOnZeroSupply() external {
+    assertEq(erc4626Upgradeable.totalAssets(), 0);
+  }
+
   // ### DEPOSIT TESTS ###
   function test_depositATokens(uint128 assets, address receiver) public {
     _validateReceiver(receiver);
@@ -74,6 +78,7 @@ contract ERC4626StataTokenUpgradeableTest is TestnetProcedures {
     assertEq(erc4626Upgradeable.balanceOf(receiver), shares);
     assertEq(IERC20(aToken).balanceOf(address(erc4626Upgradeable)), env.amount);
     assertEq(IERC20(aToken).balanceOf(user), 0);
+    assertEq(erc4626Upgradeable.totalAssets(), env.amount);
   }
 
   function test_depositATokens_self() external {