Budget.t.sol

// SPDX-License-Identifier: UNLICENSED
pragma solidity 0.8.16;

import {FirmTest} from "../../common/test/lib/FirmTest.sol";
import {RolesStub} from "../../common/test/mocks/RolesStub.sol";
import {roleFlag} from "../../common/test/mocks/RolesAuthMock.sol";
import {AvatarStub} from "../../common/test/mocks/AvatarStub.sol";
import {UpgradeableModuleProxyFactory} from "../../factory/UpgradeableModuleProxyFactory.sol";

import {TimeShift, DateTimeLib} from "../../budget/TimeShiftLib.sol";
import {SafeAware} from "../../bases/SafeAware.sol";
import "../Budget.sol";

contract BudgetTest is FirmTest {
    AvatarStub avatar;
    RolesStub roles;
    Budget budget;

    address SPENDER = account("spender");
    address RECEIVER = account("receiver");
    address SOMEONE_ELSE = account("someone else");

    function setUp() public {
        avatar = new AvatarStub();
        roles = new RolesStub();
        budget = Budget(createProxy(new Budget(), abi.encodeCall(Budget.initialize, (avatar, roles, address(0)))));
    }

    function testInitialState() public {
        assertEq(address(budget.avatar()), address(avatar));
        assertEq(address(budget.target()), address(avatar));
        assertEq(address(budget.roles()), address(roles));
    }

    function testCannotReinit() public {
        vm.expectRevert(abi.encodeWithSelector(SafeAware.AlreadyInitialized.selector));
        budget.initialize(avatar, roles, address(0));
    }

    function testCreateAllowance() public returns (uint256 allowanceId) {
        vm.prank(address(avatar));
        vm.warp(0);
        allowanceId = createDailyAllowance(SPENDER, 1);
        (
            uint256 parentId,
            uint256 amount,
            uint256 spent,
            address token,
            uint40 nextResetTime,
            address spender,
            EncodedTimeShift recurrency,
            bool isDisabled
        ) = budget.allowances(allowanceId);

        assertEq(parentId, NO_PARENT_ID);
        assertEq(amount, 10);
        assertEq(spent, 0);
        assertEq(token, address(0));
        assertEq(nextResetTime, 1 days);
        assertEq(spender, SPENDER);
        assertEq(
            bytes32(EncodedTimeShift.unwrap(recurrency)),
            bytes32(EncodedTimeShift.unwrap(TimeShift(TimeShiftLib.TimeUnit.Daily, 0).encode()))
        );
        assertFalse(isDisabled);
    }

    function testUpdateAllowanceParams() public {
        uint256 allowanceId = testCreateAllowance();

        vm.startPrank(address(avatar));
        budget.setAllowanceSpender(allowanceId, RECEIVER);
        budget.setAllowanceAmount(allowanceId, 1);
        budget.setAllowanceName(allowanceId, "new name");

        (, uint256 amount,,, uint40 nextResetTime, address spender,,) = budget.allowances(allowanceId);

        assertEq(amount, 1);
        assertEq(nextResetTime, 1 days);
        assertEq(spender, RECEIVER);
    }

    function testNotOwnerCannotCreateTopLevelAllowance() public {
        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, 0));
        createDailyAllowance(SPENDER, 0);
    }

    function testBadTimeshiftsRevert() public {
        vm.prank(address(avatar));
        vm.expectRevert(abi.encodeWithSelector(TimeShiftLib.InvalidTimeShift.selector));
        budget.createAllowance(
            NO_PARENT_ID, SPENDER, address(0), 10, TimeShift(TimeShiftLib.TimeUnit.Inherit, 0).encode(), ""
        );
    }

    function testInvalidSpenderReverts() public {
        uint8 badRoleId = 101; // RolesStub returns false to roleExists when id > 100
        vm.prank(address(avatar));
        vm.expectRevert(abi.encodeWithSelector(RolesAuth.UnexistentRole.selector, badRoleId));
        budget.createAllowance(
            NO_PARENT_ID, roleFlag(badRoleId), address(0), 10, TimeShift(TimeShiftLib.TimeUnit.Daily, 0).encode(), ""
        );
    }

    function testAllowanceIsKeptTrackOfOnSingle() public {
        uint40 initialTime = uint40(DateTimeLib.timestampFromDateTime(2022, 1, 1, 0, 0, 0));
        uint256 allowanceId = 1;

        vm.prank(address(avatar));
        vm.warp(initialTime);
        createDailyAllowance(SPENDER, allowanceId);

        assertExecutePayment(SPENDER, allowanceId, RECEIVER, 7, initialTime + 1 days);

        vm.warp(initialTime + 1 days);
        assertExecutePayment(SPENDER, allowanceId, RECEIVER, 7, initialTime + 2 days);
        assertExecutePayment(SPENDER, allowanceId, RECEIVER, 2, initialTime + 2 days);

        vm.prank(SPENDER);
        vm.expectRevert(abi.encodeWithSelector(Budget.Overbudget.selector, allowanceId, 7, 1));
        budget.executePayment(allowanceId, RECEIVER, 7, "");
    }

    function testAllowanceIsKeptTrackOfOnMulti() public {
        uint40 initialTime = uint40(DateTimeLib.timestampFromDateTime(2022, 1, 1, 0, 0, 0));
        uint256 allowanceId = 1;

        vm.prank(address(avatar));
        vm.warp(initialTime);
        createDailyAllowance(SPENDER, allowanceId);

        // max out allowance doing 5 payments of 2
        (address[] memory tos, uint256[] memory amounts) = _generateMultiPaymentArrays(5, RECEIVER, 2);
        vm.prank(SPENDER);
        budget.executeMultiPayment(allowanceId, tos, amounts, "");

        vm.prank(SPENDER);
        vm.expectRevert(abi.encodeWithSelector(Budget.Overbudget.selector, allowanceId, 1, 0));
        budget.executePayment(allowanceId, RECEIVER, 1, "");

        vm.warp(initialTime + 1 days);
        // almost max out allowance doing 4 payments of 2
        (tos, amounts) = _generateMultiPaymentArrays(4, RECEIVER, 2);
        vm.prank(SPENDER);
        budget.executeMultiPayment(allowanceId, tos, amounts, "");

        // max out on second payment
        (tos, amounts) = _generateMultiPaymentArrays(2, RECEIVER, 2);
        vm.prank(SPENDER);
        vm.expectRevert(abi.encodeWithSelector(Budget.Overbudget.selector, allowanceId, 4, 2));
        budget.executeMultiPayment(allowanceId, tos, amounts, "");
    }

    function testMultipleAllowances() public {
        uint40 initialTime = uint40(DateTimeLib.timestampFromDateTime(2022, 1, 1, 0, 0, 0));

        uint256 firstAllowanceId = 1;
        uint256 secondAllowanceId = 2;

        vm.warp(initialTime);
        vm.startPrank(address(avatar));
        createDailyAllowance(SPENDER, firstAllowanceId);
        createDailyAllowance(SPENDER, secondAllowanceId);
        vm.stopPrank();

        assertExecutePayment(SPENDER, firstAllowanceId, RECEIVER, 7, initialTime + 1 days);
        assertExecutePayment(SPENDER, secondAllowanceId, RECEIVER, 7, initialTime + 1 days);

        vm.warp(initialTime + 1 days);
        assertExecutePayment(SPENDER, firstAllowanceId, RECEIVER, 7, initialTime + 2 days);
        assertExecutePayment(SPENDER, secondAllowanceId, RECEIVER, 7, initialTime + 2 days);
    }

    function testCreateSuballowance() public returns (uint256 topLevelAllowance, uint256 subAllowance) {
        uint40 initialTime = uint40(DateTimeLib.timestampFromDateTime(2022, 1, 1, 0, 0, 0));
        vm.warp(initialTime);

        vm.prank(address(avatar));
        topLevelAllowance = budget.createAllowance(
            NO_PARENT_ID, SPENDER, address(0), 10, TimeShift(TimeShiftLib.TimeUnit.Monthly, 0).encode(), ""
        );
        vm.prank(SPENDER);
        subAllowance = budget.createAllowance(
            topLevelAllowance, SOMEONE_ELSE, address(0), 5, TimeShift(TimeShiftLib.TimeUnit.Daily, 0).encode(), ""
        );

        assertExecutePayment(SOMEONE_ELSE, subAllowance, RECEIVER, 5, initialTime + 1 days);
    }

    function testNonAdminCannotUpdateAllowanceParams() public {
        (uint256 topLevelAllowanceId, uint256 subAllowanceId) = testCreateSuballowance();

        vm.startPrank(SPENDER);

        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, NO_PARENT_ID));
        budget.setAllowanceSpender(topLevelAllowanceId, RECEIVER);

        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, NO_PARENT_ID));
        budget.setAllowanceAmount(topLevelAllowanceId, 1);

        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, NO_PARENT_ID));
        budget.setAllowanceName(topLevelAllowanceId, "new name");

        vm.stopPrank();

        vm.startPrank(address(avatar));

        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, topLevelAllowanceId));
        budget.setAllowanceSpender(subAllowanceId, RECEIVER);

        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, topLevelAllowanceId));
        budget.setAllowanceAmount(subAllowanceId, 1);

        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, topLevelAllowanceId));
        budget.setAllowanceName(subAllowanceId, "new name");

        vm.stopPrank();
    }

    function testCreateSuballowanceWithInheritedRecurrency() public {
        uint40 initialTime = uint40(DateTimeLib.timestampFromDateTime(2022, 1, 1, 0, 0, 0));
        vm.warp(initialTime);

        vm.prank(address(avatar));
        uint256 topLevelAllowance = budget.createAllowance(
            NO_PARENT_ID, SPENDER, address(0), 10, TimeShift(TimeShiftLib.TimeUnit.Daily, 0).encode(), ""
        );
        vm.prank(SPENDER);
        uint256 subAllowance = budget.createAllowance(
            topLevelAllowance, SOMEONE_ELSE, address(0), 5, TimeShift(TimeShiftLib.TimeUnit.Inherit, 0).encode(), ""
        );

        assertExecutePayment(SOMEONE_ELSE, subAllowance, RECEIVER, 5, initialTime + 1 days);
    }

    function testAllowanceChain() public {
        uint40 initialTime = uint40(DateTimeLib.timestampFromDateTime(2022, 1, 1, 0, 0, 0));
        vm.warp(initialTime);

        vm.prank(address(avatar));
        uint256 allowance1 = budget.createAllowance(
            NO_PARENT_ID, SPENDER, address(0), 10, TimeShift(TimeShiftLib.TimeUnit.Monthly, 0).encode(), ""
        );
        vm.startPrank(SPENDER);
        uint256 allowance2 = budget.createAllowance(
            allowance1, SPENDER, address(0), 5, TimeShift(TimeShiftLib.TimeUnit.Inherit, 0).encode(), ""
        );
        uint256 allowance3 = budget.createAllowance(
            allowance2, SPENDER, address(0), 2, TimeShift(TimeShiftLib.TimeUnit.Daily, 0).encode(), ""
        );
        uint256 allowance4 = budget.createAllowance(
            allowance3, SPENDER, address(0), 1, TimeShift(TimeShiftLib.TimeUnit.Inherit, 0).encode(), ""
        );
        vm.stopPrank();

        assertExecutePayment(SPENDER, allowance4, RECEIVER, 1, initialTime + 1 days);

        vm.warp(initialTime + 1 days);
        assertExecutePayment(SPENDER, allowance3, RECEIVER, 2, initialTime + 2 days);

        (,, uint256 spent1,,,,,) = budget.allowances(allowance1);
        (,, uint256 spent2,,,,,) = budget.allowances(allowance2);
        (,, uint256 spent3,,,,,) = budget.allowances(allowance3);
        (,, uint256 spent4,,,,,) = budget.allowances(allowance4);

        assertEq(spent1, 3);
        assertEq(spent2, 3);
        assertEq(spent3, 2);
        assertEq(spent4, 1); // It's one because the state doesn't get reset until a payment involving this allowance

        vm.expectRevert(abi.encodeWithSelector(Budget.Overbudget.selector, allowance3, 1, 0));
        vm.prank(SPENDER);
        budget.executePayment(allowance4, SPENDER, 1, "");
    }

    function testDisablingAllowanceBreaksChain() public {
        uint256 topLevelAllowanceId = 1;
        uint256 childAllowanceId = 4;

        testAllowanceChain(); // sets up the chain

        vm.prank(address(avatar));
        budget.setAllowanceState(topLevelAllowanceId, false);

        vm.prank(SPENDER);
        vm.expectRevert(abi.encodeWithSelector(Budget.DisabledAllowance.selector, topLevelAllowanceId));
        budget.executePayment(childAllowanceId, RECEIVER, 1, "");
    }

    function testOnlyParentAdminCanDisableAllowance() public {
        uint256 topLevelAllowanceId = 1;
        uint256 childAllowanceId = 4;

        testAllowanceChain(); // sets up the chain

        vm.prank(SPENDER);
        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, 0));
        budget.setAllowanceState(topLevelAllowanceId, false);

        vm.prank(address(avatar));
        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedNotAllowanceAdmin.selector, childAllowanceId - 1));
        budget.setAllowanceState(childAllowanceId, false);
    }

    function testCantExecuteIfNotAuthorized() public {
        vm.prank(address(avatar));
        uint256 allowanceId = 1;
        createDailyAllowance(SPENDER, allowanceId);

        vm.prank(RECEIVER);
        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedPaymentExecution.selector, allowanceId, RECEIVER));
        budget.executePayment(allowanceId, RECEIVER, 7, "");
    }

    function testAllowanceSpenderWithRoleFlags() public {
        uint256 allowanceId = 1;
        uint8 roleId = 1;
        vm.prank(address(avatar));
        createDailyAllowance(roleFlag(roleId), allowanceId);

        vm.startPrank(SPENDER);
        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedPaymentExecution.selector, allowanceId, SPENDER));
        budget.executePayment(allowanceId, RECEIVER, 7, ""); // execution fails since SPENDER doesn't have the required role yet

        roles.setRole(SPENDER, roleId, true);
        budget.executePayment(allowanceId, RECEIVER, 7, ""); // as soon as it gets the role, the payment executes
    }

    function testCantExecuteInexistentAllowance() public {
        vm.prank(SPENDER);
        vm.expectRevert(abi.encodeWithSelector(Budget.UnexistentAllowance.selector, 0));
        budget.executePayment(0, RECEIVER, 7, "");
    }

    function testCantExecuteMultiIfNotAuthorized() public {
        vm.prank(address(avatar));
        uint256 allowanceId = 1;
        createDailyAllowance(SPENDER, allowanceId);

        vm.prank(RECEIVER);
        vm.expectRevert(abi.encodeWithSelector(Budget.UnauthorizedPaymentExecution.selector, allowanceId, RECEIVER));
        (address[] memory tos, uint256[] memory amounts) = _generateMultiPaymentArrays(2, RECEIVER, 7);
        budget.executeMultiPayment(allowanceId, tos, amounts, "");
    }

    function testRevertOnBadInputToMultiPayment() public {
        vm.prank(address(avatar));
        uint256 allowanceId = 1;
        createDailyAllowance(SPENDER, allowanceId);

        vm.prank(SPENDER);
        vm.expectRevert(abi.encodeWithSelector(Budget.BadInput.selector));
        (address[] memory tos, uint256[] memory amounts) = _generateMultiPaymentArrays(2, RECEIVER, 7);
        tos = new address[](1);
        budget.executeMultiPayment(allowanceId, tos, amounts, "");
    }

    function createDailyAllowance(address spender, uint256 expectedId) public returns (uint256 allowanceId) {
        allowanceId = budget.createAllowance(
            NO_PARENT_ID, spender, address(0), 10, TimeShift(TimeShiftLib.TimeUnit.Daily, 0).encode(), ""
        );
        assertEq(allowanceId, expectedId);
    }

    function _generateMultiPaymentArrays(uint256 num, address to, uint256 amount)
        internal
        pure
        returns (address[] memory tos, uint256[] memory amounts)
    {
        tos = new address[](num);
        amounts = new uint256[](num);
        for (uint256 i = 0; i < num; i++) {
            tos[i] = to;
            amounts[i] = amount;
        }
    }

    event PaymentExecuted(
        uint256 indexed allowanceId,
        address indexed actor,
        address token,
        address indexed to,
        uint256 amount,
        uint40 nextResetTime,
        string descrition
    );

    function assertExecutePayment(
        address actor,
        uint256 allowanceId,
        address to,
        uint256 amount,
        uint40 expectedNextResetTime
    ) public {
        (,, uint256 initialSpent, address token, uint40 initialNextReset,, EncodedTimeShift shift,) =
            budget.allowances(allowanceId);

        if (block.timestamp >= initialNextReset) {
            initialSpent = 0;
        }

        vm.prank(actor);
        vm.expectEmit(true, true, true, true);
        emit PaymentExecuted(allowanceId, actor, token, to, amount, expectedNextResetTime, "");
        budget.executePayment(allowanceId, to, amount, "");

        (,, uint256 spent,, uint40 nextResetTime,,,) = budget.allowances(allowanceId);

        assertEq(spent, initialSpent + amount);
        assertEq(nextResetTime, shift.isInherited() ? 0 : expectedNextResetTime);
    }
}