From 157345110ec3d22e32687d0f6812f9f33cd9f4e7 Mon Sep 17 00:00:00 2001
From: Ben DiFrancesco <ben@scopelift.co>
Date: Thu, 21 Mar 2024 18:45:34 -0400
Subject: [PATCH] In UniStaker tests, remember depositors and surrogates and
 prevent overlap

Due to a recent, undocumented change in the Foundry fuzzer implementation, we are now
seeing test failures where the address of a yet-to-be-deployed DelegationSurrogate is
selected as a depositor. This breaks our tests. This is occuring despite the include_storage
flag being set to false, which previously was sufficient for preventing this edge case.

To accomodate this, we implement a test helper that remembers each depositor and surrogate as
they're chosen and deployed respectively, and assures there is no overlap between the two
groups of addresses.
---
 test/UniStaker.t.sol | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol
index 89124f8..23e427e 100644
--- a/test/UniStaker.t.sol
+++ b/test/UniStaker.t.sol
@@ -28,6 +28,9 @@ contract UniStakerTest is Test, PercentAssertions {
   event RewardNotifierSet(address indexed account, bool isEnabled);
   event AdminSet(address indexed oldAdmin, address indexed newAdmin);
 
+  mapping(DelegationSurrogate surrogate => bool isKnown) isKnownSurrogate;
+  mapping(address depositor => bool isKnown) isKnownDepositor;
+
   function setUp() public {
     // Set the block timestamp to an arbitrary value to avoid introducing assumptions into tests
     // based on a starting timestamp of 0, which is the default.
@@ -75,6 +78,20 @@ contract UniStakerTest is Test, PercentAssertions {
     _boundedStakeAmount = bound(_stakeAmount, 0.1e18, 25_000_000e18);
   }
 
+  // Remember each depositor and surrogate (as they're deployed) and ensure that there is
+  // no overlap between them. This is to prevent the fuzzer from selecting a surrogate as a
+  // depositor or vice versa.
+  function _assumeSafeDepositorAndSurrogate(address _depositor, address _delegatee) internal {
+    DelegationSurrogate _surrogate = uniStaker.surrogates(_delegatee);
+    isKnownDepositor[_depositor] = true;
+    isKnownSurrogate[_surrogate] = true;
+
+    vm.assume(
+      (!isKnownSurrogate[DelegationSurrogate(_depositor)])
+        && (!isKnownDepositor[address(_surrogate)])
+    );
+  }
+
   function _stake(address _depositor, uint256 _amount, address _delegatee)
     internal
     returns (UniStaker.DepositIdentifier _depositId)
@@ -85,6 +102,9 @@ contract UniStakerTest is Test, PercentAssertions {
     govToken.approve(address(uniStaker), _amount);
     _depositId = uniStaker.stake(_amount, _delegatee);
     vm.stopPrank();
+
+    // Called after the stake so the surrogate will exist
+    _assumeSafeDepositorAndSurrogate(_depositor, _delegatee);
   }
 
   function _stake(address _depositor, uint256 _amount, address _delegatee, address _beneficiary)
@@ -97,6 +117,9 @@ contract UniStakerTest is Test, PercentAssertions {
     govToken.approve(address(uniStaker), _amount);
     _depositId = uniStaker.stake(_amount, _delegatee, _beneficiary);
     vm.stopPrank();
+
+    // Called after the stake so the surrogate will exist
+    _assumeSafeDepositorAndSurrogate(_depositor, _delegatee);
   }
 
   function _fetchDeposit(UniStaker.DepositIdentifier _depositId)