diff --git a/foundry_tests/tests/Invariants.t.sol b/foundry_tests/tests/Invariants.t.sol index 44c8446a..1c4cf968 100644 --- a/foundry_tests/tests/Invariants.t.sol +++ b/foundry_tests/tests/Invariants.t.sol @@ -66,9 +66,7 @@ contract VaultInvariantTest is BaseInvariant { } function invariant_balanceAndTotalAssets() public { - if (!vaultHandler.unreported()) { - assert_balanceAndTotalAssets(); - } + assert_balanceAndTotalAssets(vaultHandler.unreported()); } function invariant_totalDebt() public { diff --git a/foundry_tests/utils/BaseInvariant.sol b/foundry_tests/utils/BaseInvariant.sol index 67178fa1..4bf3adb2 100644 --- a/foundry_tests/utils/BaseInvariant.sol +++ b/foundry_tests/utils/BaseInvariant.sol @@ -24,7 +24,7 @@ abstract contract BaseInvariant is Setup { function assert_maxWithdraw(bool unreportedLoss) public { if (unreportedLoss) { - // withdraw would revert with unreported loss so maxWithdraw is 0 + // withdraw would revert with unreported loss so maxWithdraw is totalIdle assertLe(vault.maxWithdraw(msg.sender), vault.totalIdle()); } else { assertLe(vault.maxWithdraw(msg.sender), vault.totalAssets()); @@ -124,11 +124,14 @@ abstract contract BaseInvariant is Setup { ); } - function assert_balanceAndTotalAssets() public { - assertLe( - vault.totalAssets(), - asset.balanceOf(address(strategy)) + asset.balanceOf(address(vault)) - ); + function assert_balanceAndTotalAssets(bool unreported) public { + if (!unreported) { + assertLe( + vault.totalAssets(), + asset.balanceOf(address(strategy)) + + asset.balanceOf(address(vault)) + ); + } assertEq(vault.totalIdle(), asset.balanceOf(address(vault))); }