From 687022da9ae2edb083c3a0d6c65af9f4af78edc0 Mon Sep 17 00:00:00 2001
From: Lukas <lukasstrassel@googlemail.com>
Date: Fri, 6 Sep 2024 10:47:20 +0200
Subject: [PATCH] fix: total assets (#17)

* fix: calculate totalAssets

* test: add test to ensure doesn't revert on zero
---
 .../static-a-token/ERC4626StataTokenUpgradeable.sol          | 5 +++++
 .../static-a-token/ERC4626StataTokenUpgradeable.t.sol        | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/src/periphery/contracts/static-a-token/ERC4626StataTokenUpgradeable.sol b/src/periphery/contracts/static-a-token/ERC4626StataTokenUpgradeable.sol
index 4f1b2d56..d6be297a 100644
--- 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
index e975def0..3dbe4fcd 100644
--- 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 {