From cb21fd185a846c7db986e3b84d496540fb1ec64b Mon Sep 17 00:00:00 2001
From: Schlagonia <schlagonia@gmail.com>
Date: Sat, 28 Oct 2023 14:55:58 -0600
Subject: [PATCH] feat: auto health check

---
 src/HealthCheck/BaseHealthCheck.sol           | 23 +++++++--
 .../{HealthCheck.sol => HealthCheck.t.sol}    | 51 -------------------
 src/test/mocks/MockHealthCheck.sol            | 36 +------------
 3 files changed, 21 insertions(+), 89 deletions(-)
 rename src/test/{HealthCheck.sol => HealthCheck.t.sol} (87%)

diff --git a/src/HealthCheck/BaseHealthCheck.sol b/src/HealthCheck/BaseHealthCheck.sol
index e1b5bc6..23b2090 100644
--- a/src/HealthCheck/BaseHealthCheck.sol
+++ b/src/HealthCheck/BaseHealthCheck.sol
@@ -13,9 +13,9 @@ import {BaseStrategy, ERC20} from "@tokenized-strategy/BaseStrategy.sol";
  *   `checkHealth` modifier.
  *
  *   A strategist simply needs to inherit this contract. Set
- *   the limit ratios to the desired amounts and then call
- *   `_executeHealthCheck(...)` during the  `_harvestAndReport()`
- *   execution. If the profit or loss that would be recorded is
+ *   the limit ratios to the desired amounts and then
+ *   override `_harvestAndReport()` just as they otherwise
+ *  would. If the profit or loss that would be recorded is
  *   outside the acceptable bounds the tx will revert.
  *
  *   The healthcheck does not prevent a strategy from reporting
@@ -109,6 +109,23 @@ abstract contract BaseHealthCheck is BaseStrategy {
         doHealthCheck = _doHealthCheck;
     }
 
+    /**
+     * @notice OVerrides the default {harvestAndReport} to include a healthcheck.
+     * @return _totalAssets New totalAssets post report.
+     */
+    function harvestAndReport()
+        external
+        override
+        onlySelf
+        returns (uint256 _totalAssets)
+    {
+        // Let the strategy report.
+        _totalAssets = _harvestAndReport();
+
+        // Run the healthcheck on the amount returned.
+        _executeHealthCheck(_totalAssets);
+    }
+
     /**
      * @dev To be called during a report to make sure the profit
      * or loss being recorded is within the acceptable bound.
diff --git a/src/test/HealthCheck.sol b/src/test/HealthCheck.t.sol
similarity index 87%
rename from src/test/HealthCheck.sol
rename to src/test/HealthCheck.t.sol
index 78b1c9e..56daae6 100644
--- a/src/test/HealthCheck.sol
+++ b/src/test/HealthCheck.t.sol
@@ -428,55 +428,4 @@ contract HealthCheckTest is Setup {
             "doHealthCheck should be true"
         );
     }
-
-    function test__checkHealthModifier(uint256 _amount) public {
-        vm.assume(_amount >= minFuzzAmount && _amount <= maxFuzzAmount);
-
-        // deposit
-        mintAndDepositIntoStrategy(
-            IStrategy(address(healthCheck)),
-            user,
-            _amount
-        );
-
-        assertEq(healthCheck.healthy(), true);
-        assertEq(healthCheck.availableDepositLimit(user), type(uint256).max);
-        assertEq(healthCheck.availableWithdrawLimit(user), type(uint256).max);
-
-        healthCheck.setHealthy(false);
-
-        assertEq(healthCheck.healthy(), false);
-        assertEq(healthCheck.availableDepositLimit(user), 0);
-        assertEq(healthCheck.availableWithdrawLimit(user), 0);
-
-        vm.expectRevert("unhealthy");
-        vm.prank(keeper);
-        healthCheck.report();
-
-        healthCheck.setHealthy(true);
-
-        assertEq(healthCheck.healthy(), true);
-        assertEq(healthCheck.availableDepositLimit(user), type(uint256).max);
-        assertEq(healthCheck.availableWithdrawLimit(user), type(uint256).max);
-
-        vm.prank(keeper);
-        (uint256 realProfit, ) = healthCheck.report();
-
-        // Make sure we reported the correct profit
-        assertEq(0, realProfit, "Reported profit mismatch");
-
-        // Health Check should still be on
-        assertEq(
-            healthCheck.doHealthCheck(),
-            true,
-            "doHealthCheck should be true"
-        );
-
-        skip(healthCheck.profitMaxUnlockTime());
-
-        vm.prank(user);
-        healthCheck.redeem(_amount, user, user);
-
-        assertGe(asset.balanceOf(user), _amount);
-    }
 }
diff --git a/src/test/mocks/MockHealthCheck.sol b/src/test/mocks/MockHealthCheck.sol
index a37faa0..5f4345a 100644
--- a/src/test/mocks/MockHealthCheck.sol
+++ b/src/test/mocks/MockHealthCheck.sol
@@ -22,44 +22,10 @@ contract MockHealthCheck is BaseHealthCheck {
         override
         returns (uint256 _totalAssets)
     {
-        require(_healthy(), "unhealthy");
-
         _totalAssets = asset.balanceOf(address(this));
-
-        _executeHealthCheck(_totalAssets);
-    }
-
-    // Can't deposit if its not healthy
-    function availableDepositLimit(
-        address _owner
-    ) public view override returns (uint256) {
-        if (!_healthy()) return 0;
-
-        return super.availableDepositLimit(_owner);
-    }
-
-    // Can't Withdraw if not healthy.
-    function availableWithdrawLimit(
-        address _owner
-    ) public view override returns (uint256) {
-        if (!_healthy()) return 0;
-
-        return super.availableWithdrawLimit(_owner);
-    }
-
-    function _healthy() internal view returns (bool) {
-        return healthy;
-    }
-
-    function setHealthy(bool _health) external {
-        healthy = _health;
     }
 }
 
 import {IBaseHealthCheck} from "../../HealthCheck/IBaseHealthCheck.sol";
 
-interface IMockHealthCheck is IBaseHealthCheck {
-    function healthy() external view returns (bool);
-
-    function setHealthy(bool _health) external;
-}
+interface IMockHealthCheck is IBaseHealthCheck {}