test: explicitly calculate expected withdrawable shares to account fo…

…r rounding
Layr-Labs · Feb 19, 2025 · f74dd1d · f74dd1d
1 parent 5ebed67
commit f74dd1d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/src/test/integration/IntegrationBase.t.sol b/src/test/integration/IntegrationBase.t.sol
@@ -2759,6 +2759,22 @@ abstract contract IntegrationBase is IntegrationDeployer, TypeImporter {
         return shares;
     }
 
+    function _getExpectedDSFUndelegate(User staker) internal view returns (uint expectedDepositScalingFactor) {
+        return WAD.divWad(_getBeaconChainSlashingFactor(staker));
+    }
+
+    function _getExpectedWithdrawableSharesUndelegate(User staker, IStrategy[] memory strategies, uint[] memory shares) internal view returns (uint[] memory){
+        uint[] memory expectedWithdrawableShares = new uint[](strategies.length);
+        for (uint i = 0; i < strategies.length; i++) {
+            if (strategies[i] == BEACONCHAIN_ETH_STRAT) {
+                expectedWithdrawableShares[i] = shares[i].mulWad(_getExpectedDSFUndelegate(staker)).mulWad(_getBeaconChainSlashingFactor(staker));
+            } else {
+                expectedWithdrawableShares[i] = shares[i];
+             }
+        }
+        return expectedWithdrawableShares;
+    }
+
     function _getPrevWithdrawableShares(User staker, IStrategy[] memory strategies) internal timewarp() returns (uint[] memory) {
         return _getWithdrawableShares(staker, strategies);
     }

diff --git a/src/test/integration/IntegrationChecks.t.sol b/src/test/integration/IntegrationChecks.t.sol
@@ -358,7 +358,8 @@ contract IntegrationCheckUtils is IntegrationBase {
         assert_Snap_Unchanged_TokenBalances(staker, "staker should not have any change in underlying token balances");
         assert_Snap_Unchanged_TokenBalances(operator, "operator should not have any change in underlying token balances");
         assert_Snap_Added_Staker_DepositShares(staker, strategies, shares, "staker should have received expected deposit shares");
-        assert_Snap_Added_Staker_WithdrawableShares(staker, strategies, shares, "staker should have received expected withdrawable shares");
+        uint[] memory expectedWithdrawableShares = _getExpectedWithdrawableSharesUndelegate(staker, strategies, shares);
+        assert_Snap_Added_Staker_WithdrawableShares(staker, strategies, expectedWithdrawableShares, "staker should have received expected withdrawable shares");
         assert_Snap_Unchanged_OperatorShares(operator, "operator should have shares unchanged");
         assert_Snap_Unchanged_StrategyShares(strategies, "strategies should have total shares unchanged");
     }