Nested Pranking for Foundry Tests

BaseTest.sol

// SPDX-License-Identifier: MIT

pragma solidity >=0.6.12;
pragma experimental ABIEncoderV2;

import "forge-std/Test.sol";

/**
 * Base testing contract with nested pranking capabilities.
 * @author Dalton Sweeney
 */
abstract contract BaseTest is Test {
    struct Prank {
        address sender;
        address origin;
    }

    /// @notice Stack of active pranks
    Prank[] private _pranks;

    /// @notice Start a prank where msg.sender == `sender`. The active prank will be paused and then resumed
    /// on the next call to _stopPrank().
    /// @param sender address to prank msg.sender
    function _startPrank(address sender) internal {
        _startPrank(sender, address(0));
    }

    /// @notice Start a prank where msg.sender == `sender` and tx.origin == `origin`. If you pass the zero address
    /// for origin, then tx.origin IS NOT pranked. The active prank will be paused and then resumed on the next call
    /// to _stopPrank()
    /// @param sender address to prank msg.sender
    /// @param origin address to prank tx.origin. Use the zero address to indicate you do not want to prank tx.origin
    function _startPrank(address sender, address origin) internal {
        // Pause the active prank
        if (_pranks.length > 0) {
            vm.stopPrank();
        }

        _pranks.push(Prank({sender: sender, origin: origin}));

        // 0 address means don't prank tx.origin
        if (origin == address(0)) {
            vm.startPrank(sender);
        } else {
            vm.startPrank(sender, origin);
        }
    }

    /// @notice Stop the current prank. Resumes the previous prank if it exists
    function _stopPrank() internal {
        require(_pranks.length > 0, "No active pranks");

        vm.stopPrank();
        _pranks.pop();

        // Resume the previous prank
        if (_pranks.length > 0) {
            Prank memory prank = _pranks[_pranks.length - 1];
            if (prank.origin == address(0)) {
                vm.startPrank(prank.sender);
            } else {
                vm.startPrank(prank.sender, prank.origin);
            }
        }
    }

    modifier prankingSender(address sender) {
        _startPrank(sender);
        _;
        _stopPrank();
    }

    modifier prankingSenderAndOrigin(address sender, address origin) {
        _startPrank(sender, origin);
        _;
        _stopPrank();
    }
}