From 42802408cef99cf75d479e418f63ab0d30e21974 Mon Sep 17 00:00:00 2001
From: Matthew Brown <mattfdbrown@gmail.com>
Date: Thu, 6 Jun 2024 18:11:38 -0500
Subject: [PATCH] get the totalAssets at a timestamp

---
 contracts/test/MockERC4626.sol | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/contracts/test/MockERC4626.sol b/contracts/test/MockERC4626.sol
index ec0b3d2e4..7168f0aaf 100644
--- a/contracts/test/MockERC4626.sol
+++ b/contracts/test/MockERC4626.sol
@@ -150,7 +150,13 @@ contract MockERC4626 is ERC4626, MultiRolesAuthority {
     }
 
     function totalAssets() public view override returns (uint256) {
-        return asset.balanceOf(address(this)) + _getAccruedInterest();
+        return
+            asset.balanceOf(address(this)) +
+            _getAccruedInterest(block.timestamp);
+    }
+
+    function totalAssets(uint256 timestamp) public view returns (uint256) {
+        return asset.balanceOf(address(this)) + _getAccruedInterest(timestamp);
     }
 
     /// Mock ///
@@ -178,14 +184,16 @@ contract MockERC4626 is ERC4626, MultiRolesAuthority {
     }
 
     function _accrue() internal {
-        uint256 interest = _getAccruedInterest();
+        uint256 interest = _getAccruedInterest(block.timestamp);
         if (interest > 0) {
             ERC20Mintable(address(asset)).mint(interest);
         }
         _lastUpdated = block.timestamp;
     }
 
-    function _getAccruedInterest() internal view returns (uint256) {
+    function _getAccruedInterest(
+        uint256 timestamp
+    ) internal view returns (uint256) {
         // If the rate is zero, no interest has accrued.
         if (_rate == 0) {
             return 0;
@@ -194,14 +202,12 @@ contract MockERC4626 is ERC4626, MultiRolesAuthority {
         // If the block timestamp is less than last updated, the accrual
         // calculation will underflow. This can occur when using anvil's state
         // snapshots.
-        if (block.timestamp < _lastUpdated) {
+        if (timestamp < _lastUpdated) {
             return 0;
         }
 
         // base_balance = base_balance * (1 + r * t)
-        uint256 timeElapsed = (block.timestamp - _lastUpdated).divDown(
-            365 days
-        );
+        uint256 timeElapsed = (timestamp - _lastUpdated).divDown(365 days);
         uint256 accrued = asset.balanceOf(address(this)).mulDown(
             _rate.mulDown(timeElapsed)
         );