gru_emoney_token-factory/test/upgrade/UpgradeTest.t.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "forge-std/Test.sol";
import "../../src/eMoneyToken.sol";
import "../../src/PolicyManager.sol";
import "../../src/ComplianceRegistry.sol";
import "../../src/DebtRegistry.sol";
import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";

/**
 * @title UpgradeTest
 * @notice Tests UUPS upgrade functionality and storage layout compatibility
 */
contract UpgradeTest is Test {
    eMoneyToken public tokenProxy;
    eMoneyToken public implementationV1;
    eMoneyToken public implementationV2;
    PolicyManager public policyManager;
    ComplianceRegistry public complianceRegistry;
    DebtRegistry public debtRegistry;

    address public admin;
    address public issuer;
    address public user1;

    function setUp() public {
        admin = address(0x1);
        issuer = address(0x2);
        user1 = address(0x10);

        complianceRegistry = new ComplianceRegistry(admin);
        debtRegistry = new DebtRegistry(admin);
        policyManager = new PolicyManager(admin, address(complianceRegistry), address(debtRegistry));

        // Deploy V1 implementation
        implementationV1 = new eMoneyToken();

        // Deploy proxy with V1
        bytes memory initData = abi.encodeWithSelector(
            eMoneyToken.initialize.selector,
            "Test Token",
            "TEST",
            18,
            issuer,
            address(policyManager),
            address(debtRegistry),
            address(complianceRegistry)
        );
        ERC1967Proxy proxy = new ERC1967Proxy(address(implementationV1), initData);
        tokenProxy = eMoneyToken(address(proxy));

        vm.startPrank(issuer);
        tokenProxy.grantRole(tokenProxy.ISSUER_ROLE(), issuer);
        vm.stopPrank();

        vm.startPrank(admin);
        complianceRegistry.grantRole(complianceRegistry.COMPLIANCE_ROLE(), admin);
        complianceRegistry.setCompliance(user1, true, 1, bytes32(0));
        policyManager.grantRole(policyManager.POLICY_OPERATOR_ROLE(), admin);
        policyManager.setLienMode(address(tokenProxy), 2);
        vm.stopPrank();
    }

    /**
     * @notice V2 implementation with same storage layout
     */
    function deployV2Implementation() internal returns (eMoneyToken) {
        return new eMoneyToken();
    }

    function test_upgrade_preservesState() public {
        // Set up some state
        vm.prank(issuer);
        tokenProxy.mint(user1, 1000, bytes32(0));

        uint256 balanceBefore = tokenProxy.balanceOf(user1);
        string memory nameBefore = tokenProxy.name();
        string memory symbolBefore = tokenProxy.symbol();
        uint8 decimalsBefore = tokenProxy.decimals();

        // Deploy V2
        implementationV2 = deployV2Implementation();

        // Upgrade - OpenZeppelin v5 uses upgradeToAndCall
        vm.prank(issuer);
        UUPSUpgradeable(address(tokenProxy)).upgradeToAndCall(address(implementationV2), "");

        // Verify state preserved
        assertEq(tokenProxy.balanceOf(user1), balanceBefore);
        assertEq(keccak256(bytes(tokenProxy.name())), keccak256(bytes(nameBefore)));
        assertEq(keccak256(bytes(tokenProxy.symbol())), keccak256(bytes(symbolBefore)));
        assertEq(tokenProxy.decimals(), decimalsBefore);
    }

    function test_upgrade_preservesRoles() public {
        bytes32 issuerRole = tokenProxy.ISSUER_ROLE();
        bytes32 enforcementRole = tokenProxy.ENFORCEMENT_ROLE();
        bytes32 adminRole = tokenProxy.DEFAULT_ADMIN_ROLE();

        // Verify issuer has roles before upgrade
        assertTrue(tokenProxy.hasRole(issuerRole, issuer));
        assertTrue(tokenProxy.hasRole(adminRole, issuer));

        // Deploy V2 and upgrade
        implementationV2 = deployV2Implementation();
        vm.prank(issuer);
        UUPSUpgradeable(address(tokenProxy)).upgradeToAndCall(address(implementationV2), "");

        // Verify roles preserved
        assertTrue(tokenProxy.hasRole(issuerRole, issuer));
        assertTrue(tokenProxy.hasRole(adminRole, issuer));
    }

    function test_upgrade_functionalityStillWorks() public {
        vm.prank(issuer);
        tokenProxy.mint(user1, 1000, bytes32(0));

        // Deploy V2 and upgrade
        implementationV2 = deployV2Implementation();
        vm.prank(issuer);
        UUPSUpgradeable(address(tokenProxy)).upgradeToAndCall(address(implementationV2), "");

        // Test mint still works
        vm.prank(issuer);
        tokenProxy.mint(user1, 500, bytes32(0));
        assertEq(tokenProxy.balanceOf(user1), 1500);

        // Test transfer still works
        vm.prank(user1);
        tokenProxy.transfer(admin, 200);
        assertEq(tokenProxy.balanceOf(user1), 1300);
        assertEq(tokenProxy.balanceOf(admin), 200);
    }

    function test_upgrade_registryAddressesPreserved() public {
        // Deploy V2 and upgrade
        implementationV2 = deployV2Implementation();
        vm.prank(issuer);
        UUPSUpgradeable(address(tokenProxy)).upgradeToAndCall(address(implementationV2), "");

        // Verify registry addresses are still accessible
        assertEq(address(tokenProxy.policyManager()), address(policyManager));
        assertEq(address(tokenProxy.debtRegistry()), address(debtRegistry));
        assertEq(address(tokenProxy.complianceRegistry()), address(complianceRegistry));
    }

    function test_upgrade_onlyAdminCanUpgrade() public {
        implementationV2 = deployV2Implementation();

        // Non-admin cannot upgrade
        vm.prank(user1);
        vm.expectRevert();
        UUPSUpgradeable(address(tokenProxy)).upgradeToAndCall(address(implementationV2), "");

        // Admin can upgrade
        vm.prank(issuer);
        UUPSUpgradeable(address(tokenProxy)).upgradeToAndCall(address(implementationV2), "");
    }

    function test_upgrade_freeBalanceStillWorks() public {
        vm.prank(issuer);
        tokenProxy.mint(user1, 1000, bytes32(0));

        vm.startPrank(admin);
        debtRegistry.grantRole(debtRegistry.DEBT_AUTHORITY_ROLE(), admin);
        debtRegistry.placeLien(user1, 300, 0, 1, bytes32(0));
        vm.stopPrank();

        uint256 freeBalanceBefore = tokenProxy.freeBalanceOf(user1);
        assertEq(freeBalanceBefore, 700);

        // Deploy V2 and upgrade
        implementationV2 = deployV2Implementation();
        vm.prank(issuer);
        UUPSUpgradeable(address(tokenProxy)).upgradeToAndCall(address(implementationV2), "");

        // Verify freeBalance still works
        uint256 freeBalanceAfter = tokenProxy.freeBalanceOf(user1);
        assertEq(freeBalanceAfter, freeBalanceBefore);
    }
}