Skip auto-funding for accounts that were explicitly dealt

crates/revive-strategy/src/cheatcodes/mock_handler.rs

@@ -5,7 +5,7 @@ use std::{
 };
 
 use alloy_primitives::{Address, Bytes, map::foldhash::HashMap, ruint::aliases::U256};
-use foundry_cheatcodes::{Ecx, MockCallDataContext, MockCallReturnData};
+use foundry_cheatcodes::{DealRecord, Ecx, MockCallDataContext, MockCallReturnData};
 use foundry_evm::constants::CHEATCODE_ADDRESS;
 use polkadot_sdk::{
     frame_system,
@@ -60,12 +60,23 @@ impl MockHandlerImpl {
         state.mocked_functions = mock_inner.mocked_functions.clone();
     }
 
+    /// Funds pranked fuzz addresses with u128::MAX so they can make calls in pallet-revive.
+    /// Skips accounts that were explicitly dealt to via vm.deal() to preserve balance assertions.
+    ///
+    /// Note: This only affects accounts that have never been dealt to. Accounts created from
+    /// within a test contract (e.g., via `self` calls) are not skipped and will be funded if
+    /// their balance is 0. This is intentional - vm.deal() is an explicit user action that
+    /// should be respected, while internal contract creations should still get auto-funding.
     pub(crate) fn fund_pranked_accounts(&self, account: Address) {
-        // Fuzzed prank addresses have no balance, so they won't exist in revive, and
-        // calls will fail, this is not a problem when running in REVM.
-        // TODO: Figure it out why this is still needed.
-        let balance = Pallet::<Runtime>::evm_balance(&H160::from_slice(account.as_slice()));
-        if balance == 0.into() {
+        let mock_inner = self.inner.borrow();
+
+        // Skip accounts that were explicitly dealt to via vm.deal()
+        if mock_inner.eth_deals.iter().any(|deal| deal.address == account) {
+            return;
+        }
+
+        let pvm_balance = Pallet::<Runtime>::evm_balance(&H160::from_slice(account.as_slice()));
+        if pvm_balance == 0.into() {
             Pallet::<Runtime>::set_evm_balance(
                 &H160::from_slice(account.as_slice()),
                 u128::MAX.into(),
@@ -193,6 +204,8 @@ struct MockHandlerInner<T: frame_system::Config + pallet_revive::Config> {
 
     pub mocked_calls: HashMap<Address, BTreeMap<MockCallDataContext, VecDeque<MockCallReturnData>>>,
     pub mocked_functions: HashMap<Address, HashMap<Bytes, Address>>,
+    /// Records of accounts that were explicitly dealt to via vm.deal().
+    pub eth_deals: Vec<DealRecord>,
 }
 
 impl MockHandlerInner<Runtime> {
@@ -222,6 +235,7 @@ impl MockHandlerInner<Runtime> {
             mocked_calls: state.mocked_calls.clone(),
             callee: callee.map(|addr| H160::from_slice(addr.as_slice())).unwrap_or_default(),
             mocked_functions: state.mocked_functions.clone(),
+            eth_deals: state.eth_deals.clone(),
         }
     }
 }

testdata/default/revive/Prank.t.sol

@@ -675,3 +675,51 @@ contract Counter {
         number++;
     }
 }
+
+contract FundPrankedAccountsReproTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    address alice = address(0xA11CE);
+
+    function test_PrankDealtAccountWithZeroBalance() public {
+        vm.deal(alice, 3 ether);
+        assertEq(alice.balance, 3 ether, "alice should have 3 ether after deal");
+
+        vm.deal(alice, 0);
+        assertEq(alice.balance, 0, "alice should have 0 balance after second deal");
+
+        vm.startPrank(alice);
+        assertEq(alice.balance, 0, "alice balance should remain 0, not be reset to u128::MAX");
+        vm.stopPrank();
+    }
+
+    function test_PrankFuzzedAddressGetsFunded() public {
+        address fuzzedAddr = address(0xF022ED);
+        uint256 initialBalance = fuzzedAddr.balance;
+
+        vm.prank(fuzzedAddr);
+
+        assertTrue(fuzzedAddr.balance >= initialBalance, "fuzzed address should be auto-funded");
+    }
+
+    function test_PrankDealtAccountWithBalance() public {
+        vm.deal(alice, 5 ether);
+        assertEq(alice.balance, 5 ether, "alice should have 5 ether");
+
+        vm.startPrank(alice);
+        assertEq(alice.balance, 5 ether, "alice balance should remain 5 ether");
+        vm.stopPrank();
+    }
+
+    function test_PrankAfterMultipleDeals() public {
+        vm.deal(alice, 10 ether);
+        vm.deal(alice, 5 ether);
+        vm.deal(alice, 2 ether);
+
+        assertEq(alice.balance, 2 ether, "alice should have 2 ether after multiple deals");
+
+        vm.startPrank(alice);
+        assertEq(alice.balance, 2 ether, "alice balance should remain 2 ether after prank");
+        vm.stopPrank();
+    }
+}