proposal-25-restore-rewards/test/ProposalUtils.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.6.12;
pragma experimental ABIEncoderV2;

import { Utils } from "./Utils.sol";
import { Proposal, IGovernance } from "./interfaces/IGovernance.sol";

import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";

contract ProposalUtils is Utils {
    IGovernance internal governance = IGovernance(payable(_governanceAddress));

    function getProposalExecutableTime(uint256 proposalId) internal view returns (uint256) {
        Proposal memory proposal = getProposal(proposalId);
        return proposal.endTime + PROPOSAL_LOCKED_DURATION + 1 seconds;
    }

    function getProposal(uint256 proposalId) internal view returns (Proposal memory) {
        return governance.proposals(proposalId);
    }

    function hasProposal(uint256 proposalId) internal view returns (bool) {
        return governance.proposalCount() >= proposalId;
    }

    function waitUntilExecutable(uint256 proposalId) internal {
        uint256 proposalExecutableTime = getProposalExecutableTime(proposalId);
        require(block.timestamp < proposalExecutableTime + PROPOSAL_EXECUTION_MAX_DURATION, "Too late to execute proposal");

        vm.warp(proposalExecutableTime);
    }

    function proposeAndVote(address proposalAddress) public returns (uint256) {
        retrieveAndLockBalance(TEST_PRIVATE_KEY_ONE, TEST_ADDRESS_ONE, PROPOSAL_THRESHOLD + 1 ether);

        /* ----------PROPOSER------------ */
        vm.startPrank(TEST_ADDRESS_ONE);

        uint256 proposalId = governance.propose(proposalAddress, PROPOSAL_DESCRIPTION);

        // TIME-TRAVEL
        vm.warp(block.timestamp + 6 hours);

        governance.castVote(proposalId, true);

        vm.stopPrank();
        /* ------------------------------ */

        return proposalId;
    }

    function proposeAndExecute(address proposalAddress) public {
        uint256 proposalId = proposeAndVote(proposalAddress);

        waitUntilExecutable(proposalId);
        IGovernance(_governanceAddress).execute(proposalId);

        returnTokensToGovernanceAfterVoting(proposalId);
    }

    function returnTokensToGovernanceAfterVoting(uint256 proposalId) public {
        uint256 retrievedAmount = PROPOSAL_THRESHOLD + 1 ether;

        uint256 proposalExecutableTime = getProposalExecutableTime(proposalId);
        uint256 tokensUnlockTime = proposalExecutableTime + 4 days + 1 seconds;

        if(block.timestamp < tokensUnlockTime) vm.warp(tokensUnlockTime);

        vm.startPrank(TEST_ADDRESS_ONE);

        governance.unlock(retrievedAmount);
        IERC20(_tokenAddress).transfer(_governanceAddress, retrievedAmount);

        vm.stopPrank();
    }
}