Skip to content

Commit

Permalink
add upgrade unit test for every upgradable contract
Browse files Browse the repository at this point in the history
  • Loading branch information
@skosito
skosito committed Oct 11, 2024
1 parent 0638977 commit ec2fd2a
Show file tree
Hide file tree
Showing 11 changed files with 1,157 additions and 111 deletions.
33 changes: 33 additions & 0 deletions v2/test/ERC20Custody.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import "forge-std/Vm.sol";
import "./utils/ReceiverEVM.sol";

import "./utils/TestERC20.sol";
import "./utils/upgrades/ERC20CustodyUpgradeTest.sol";

import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
Expand Down Expand Up @@ -39,6 +40,8 @@ contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiv
error AccessControlUnauthorizedAccount(address account, bytes32 neededRole);
error LegacyMethodsNotSupported();

event WithdrawnV2(address indexed to, address indexed token, uint256 amount);

bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE");
bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE");
bytes32 public constant ASSET_HANDLER_ROLE = keccak256("ASSET_HANDLER_ROLE");
Expand Down Expand Up @@ -552,4 +555,34 @@ contract ERC20CustodyTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiv
vm.expectRevert(NotWhitelisted.selector);
custody.deposit(abi.encodePacked(destination), token, 1000, message);
}

function testUpgradeAndWithdraw() public {
// upgrade
Upgrades.upgradeProxy(address(custody), "ERC20CustodyUpgradeTest.sol", "", owner);
ERC20CustodyUpgradeTest custodyV2 = ERC20CustodyUpgradeTest(address(custody));
// withdraw
uint256 amount = 100_000;
uint256 balanceBefore = token.balanceOf(destination);
assertEq(balanceBefore, 0);
uint256 balanceBeforeCustody = token.balanceOf(address(custodyV2));

bytes memory transferData = abi.encodeWithSignature("transfer(address,uint256)", address(destination), amount);
vm.expectCall(address(token), 0, transferData);
vm.expectEmit(true, true, true, true, address(custodyV2));
emit WithdrawnV2(destination, address(token), amount);
vm.prank(tssAddress);
custodyV2.withdraw(destination, address(token), amount);

// Verify that the tokens were transferred to the destination address
uint256 balanceAfter = token.balanceOf(destination);
assertEq(balanceAfter, amount);

// Verify that the tokens were substracted from custody
uint256 balanceAfterCustody = token.balanceOf(address(custodyV2));
assertEq(balanceAfterCustody, balanceBeforeCustody - amount);

// Verify that gateway doesn't hold any tokens
uint256 balanceGateway = token.balanceOf(address(gateway));
assertEq(balanceGateway, 0);
}
}
32 changes: 30 additions & 2 deletions v2/test/GatewayEVM.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import "forge-std/Vm.sol";
import "./utils/ReceiverEVM.sol";

import "./utils/TestERC20.sol";
import "./utils/upgrades/GatewayEVMUpgradeTest.sol";

import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
Expand All @@ -24,7 +25,6 @@ import "./utils/Zeta.non-eth.sol";
contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiverEVMEvents, IERC20CustodyEvents {
using SafeERC20 for IERC20;

address proxy;
GatewayEVM gateway;
ReceiverEVM receiver;
ERC20Custody custody;
Expand All @@ -40,6 +40,8 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver
error EnforcedPause();
error AccessControlUnauthorizedAccount(address account, bytes32 neededRole);

event ExecutedV2(address indexed destination, uint256 value, bytes data);

bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE");
bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE");
bytes32 public constant ASSET_HANDLER_ROLE = keccak256("ASSET_HANDLER_ROLE");
Expand All @@ -54,7 +56,7 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver
token = new TestERC20("test", "TTK");

zeta = new ZetaNonEth(tssAddress, tssAddress);
proxy = Upgrades.deployUUPSProxy(
address proxy = Upgrades.deployUUPSProxy(
"GatewayEVM.sol", abi.encodeCall(GatewayEVM.initialize, (tssAddress, address(zeta), owner))
);
gateway = GatewayEVM(proxy);
Expand Down Expand Up @@ -364,6 +366,32 @@ contract GatewayEVMTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiver
vm.expectRevert(ZeroAddress.selector);
gateway.executeRevert{ value: value }(address(0), data, revertContext);
}

function testUpgradeAndForwardCallToReceivePayable() public {
// upgrade
Upgrades.upgradeProxy(address(gateway), "GatewayEVMUpgradeTest.sol", "", owner);
GatewayEVMUpgradeTest gatewayUpgradeTest = GatewayEVMUpgradeTest(address(gateway));
// call
address custodyBeforeUpgrade = gateway.custody();
address tssBeforeUpgrade = gateway.tssAddress();

string memory str = "Hello, Foundry!";
uint256 num = 42;
bool flag = true;
uint256 value = 1 ether;

bytes memory data = abi.encodeWithSignature("receivePayable(string,uint256,bool)", str, num, flag);
vm.expectCall(address(receiver), value, data);
vm.expectEmit(true, true, true, true, address(receiver));
emit ReceivedPayable(address(gateway), value, str, num, flag);
vm.expectEmit(true, true, true, true, address(gateway));
emit ExecutedV2(address(receiver), value, data);
vm.prank(tssAddress);
gatewayUpgradeTest.execute{ value: value }(address(receiver), data);

assertEq(custodyBeforeUpgrade, gateway.custody());
assertEq(tssBeforeUpgrade, gateway.tssAddress());
}
}

contract GatewayEVMInboundTest is Test, IGatewayEVMErrors, IGatewayEVMEvents, IReceiverEVMEvents {
Expand Down
105 changes: 0 additions & 105 deletions v2/test/GatewayEVMUpgrade.t.sol
View file

This file was deleted.

41 changes: 41 additions & 0 deletions v2/test/GatewayZEVM.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import "../contracts/zevm/SystemContract.sol";
import "./utils/TestUniversalContract.sol";

import "./utils/WZETA.sol";
import "./utils/upgrades/GatewayZEVMUpgradeTest.sol";

import "../contracts/zevm/GatewayZEVM.sol";
import "../contracts/zevm/ZRC20.sol";
Expand All @@ -32,6 +33,19 @@ contract GatewayZEVMInboundTest is Test, IGatewayZEVMEvents, IGatewayZEVMErrors
error ZeroAddress();
error LowBalance();

event WithdrawnV2(
address indexed sender,
uint256 indexed chainId,
bytes receiver,
address zrc20,
uint256 value,
uint256 gasfee,
uint256 protocolFlatFee,
bytes message,
CallOptions callOptions,
RevertOptions revertOptions
);

function setUp() public {
owner = address(this);
addr1 = address(0x1234);
Expand Down Expand Up @@ -585,6 +599,33 @@ contract GatewayZEVMInboundTest is Test, IGatewayZEVMEvents, IGatewayZEVMErrors
emit Called(owner, address(zrc20), abi.encodePacked(addr1), message, callOptions, revertOptions);
gateway.call(abi.encodePacked(addr1), address(zrc20), message, callOptions, revertOptions);
}

function testUpgradeAndWithdrawZRC20() public {
// upgrade
Upgrades.upgradeProxy(proxy, "GatewayZEVMUpgradeTest.sol", "", owner);
GatewayZEVMUpgradeTest gatewayUpgradeTest = GatewayZEVMUpgradeTest(proxy);
// withdraw
uint256 amount = 1;
uint256 ownerBalanceBefore = zrc20.balanceOf(owner);

vm.expectEmit(true, true, true, true, address(gatewayUpgradeTest));
emit WithdrawnV2(
owner,
0,
abi.encodePacked(addr1),
address(zrc20),
amount,
0,
zrc20.PROTOCOL_FLAT_FEE(),
"",
CallOptions({ gasLimit: 0, isArbitraryCall: true }),
revertOptions
);
gatewayUpgradeTest.withdraw(abi.encodePacked(addr1), amount, address(zrc20), revertOptions);

uint256 ownerBalanceAfter = zrc20.balanceOf(owner);
assertEq(ownerBalanceBefore - amount, ownerBalanceAfter);
}
}

contract GatewayZEVMOutboundTest is Test, IGatewayZEVMEvents, IGatewayZEVMErrors {
Expand Down
23 changes: 23 additions & 0 deletions v2/test/ZetaConnectorNative.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import "forge-std/Vm.sol";
import "./utils/ReceiverEVM.sol";

import "./utils/TestERC20.sol";
import "./utils/upgrades/ZetaConnectorNativeUpgradeTest.sol";

import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
Expand Down Expand Up @@ -43,6 +44,8 @@ contract ZetaConnectorNativeTest is
error EnforcedPause();
error AccessControlUnauthorizedAccount(address account, bytes32 neededRole);

event WithdrawnV2(address indexed to, uint256 amount);

bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE");
bytes32 public constant PAUSER_ROLE = keccak256("PAUSER_ROLE");
bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE");
Expand Down Expand Up @@ -350,4 +353,24 @@ contract ZetaConnectorNativeTest is
vm.expectRevert(abi.encodeWithSelector(AccessControlUnauthorizedAccount.selector, owner, WITHDRAWER_ROLE));
zetaConnector.withdrawAndRevert(address(receiver), amount, data, internalSendHash, revertContext);
}

function testUpgradeAndWithdraw() public {
// upgrade
Upgrades.upgradeProxy(address(zetaConnector), "ZetaConnectorNativeUpgradeTest.sol", "", owner);
ZetaConnectorNativeUpgradeTest zetaConnectorV2 = ZetaConnectorNativeUpgradeTest(address(zetaConnector));
// withdraw
uint256 amount = 100_000;
bytes32 internalSendHash = "";
uint256 balanceBefore = zetaToken.balanceOf(destination);
assertEq(balanceBefore, 0);

bytes memory data = abi.encodeWithSignature("transfer(address,uint256)", destination, amount);
vm.expectCall(address(zetaToken), 0, data);
vm.expectEmit(true, true, true, true, address(zetaConnectorV2));
emit WithdrawnV2(destination, amount);
vm.prank(tssAddress);
zetaConnectorV2.withdraw(destination, amount, internalSendHash);
uint256 balanceAfter = zetaToken.balanceOf(destination);
assertEq(balanceAfter, amount);
}
}
24 changes: 24 additions & 0 deletions v2/test/ZetaConnectorNonNative.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import "forge-std/Vm.sol";
import "./utils/ReceiverEVM.sol";

import "./utils/TestERC20.sol";
import "./utils/upgrades/ZetaConnectorNonNativeUpgradeTest.sol";

import "../contracts/evm/ERC20Custody.sol";
import "../contracts/evm/GatewayEVM.sol";
Expand Down Expand Up @@ -44,6 +45,8 @@ contract ZetaConnectorNonNativeTest is
error ExceedsMaxSupply();
error EnforcedPause();

event WithdrawnV2(address indexed to, uint256 amount);

bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE");
bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE");
bytes32 public constant PAUSER_ROLE = keccak256("PAUSER_ROLE");
Expand Down Expand Up @@ -357,4 +360,25 @@ contract ZetaConnectorNonNativeTest is
vm.expectRevert(abi.encodeWithSelector(AccessControlUnauthorizedAccount.selector, owner, TSS_ROLE));
zetaConnector.setMaxSupply(10_000);
}

function testUpgradeAndWithdraw() public {
// upgrade
Upgrades.upgradeProxy(address(zetaConnector), "ZetaConnectorNonNativeUpgradeTest.sol", "", owner);
ZetaConnectorNonNativeUpgradeTest zetaConnectorV2 = ZetaConnectorNonNativeUpgradeTest(address(zetaConnector));
// withdraw
uint256 amount = 100_000;
uint256 balanceBefore = zetaToken.balanceOf(destination);
assertEq(balanceBefore, 0);
bytes32 internalSendHash = "";

bytes memory data =
abi.encodeWithSignature("mint(address,uint256,bytes32)", destination, amount, internalSendHash);
vm.expectCall(address(zetaToken), 0, data);
vm.expectEmit(true, true, true, true, address(zetaConnectorV2));
emit WithdrawnV2(destination, amount);
vm.prank(tssAddress);
zetaConnectorV2.withdraw(destination, amount, internalSendHash);
uint256 balanceAfter = zetaToken.balanceOf(destination);
assertEq(balanceAfter, amount);
}
}
Loading

0 comments on commit ec2fd2a

Please sign in to comment.