upgrade custody in upgrade tests

cmd/zetae2e/local/local.go

@@ -46,7 +46,7 @@ const (
  flagTestV2Migration = "test-v2-migration"
  flagSkipTrackerCheck = "skip-tracker-check"
  flagSkipPrecompiles = "skip-precompiles"
- flagUpgradeGateways = "upgrade-gateways"
+ flagUpgradeContracts = "upgrade-contracts"
 )
 
 var (
@@ -84,7 +84,7 @@ func NewLocalCmd() *cobra.Command {
  cmd.Flags().Bool(flagTestV2Migration, false, "set to true to run tests for v2 contracts migration test")
  cmd.Flags().Bool(flagSkipTrackerCheck, false, "set to true to skip tracker check at the end of the tests")
  cmd.Flags().Bool(flagSkipPrecompiles, false, "set to true to skip stateful precompiled contracts test")
- cmd.Flags().Bool(flagUpgradeGateways, false, "set to true to upgrade gateways during setup for ZEVM and EVM")
+ cmd.Flags().Bool(flagUpgradeContracts, false, "set to true to upgrade gateways during setup for ZEVM and EVM")
 
  return cmd
 }
@@ -114,7 +114,7 @@ func localE2ETest(cmd *cobra.Command, _ []string) {
  testV2 = must(cmd.Flags().GetBool(flagTestV2))
  testV2Migration = must(cmd.Flags().GetBool(flagTestV2Migration))
  skipPrecompiles = must(cmd.Flags().GetBool(flagSkipPrecompiles))
- upgradeGateways = must(cmd.Flags().GetBool(flagUpgradeGateways))
+ upgradeContracts = must(cmd.Flags().GetBool(flagUpgradeContracts))
  )
 
  logger := runner.NewLogger(verbose, color.FgWhite, "setup")
@@ -413,9 +413,9 @@ func localE2ETest(cmd *cobra.Command, _ []string) {
  eg.Go(tonTestRoutine(conf, deployerRunner, verbose, tonTests...))
  }
 
- // upgrade gateways
- if upgradeGateways {
- deployerRunner.UpgradeGateways()
+ // upgrade contracts
+ if upgradeContracts {
+ deployerRunner.UpgradeContracts()
  }
 
  if testV2 || testV2Migration {

contrib/localnet/orchestrator/start-zetae2e.sh

@@ -258,9 +258,9 @@ if [ "$LOCALNET_MODE" == "upgrade" ]; then
  # When the upgrade height is greater than 100 for upgrade test, the Bitcoin tests have been run once, therefore the Bitcoin wallet is already set up
  # Use light flag to skip advanced tests
  if [ "$UPGRADE_HEIGHT" -lt 100 ]; then
- zetae2e local $E2E_ARGS --skip-setup --config "$deployed_config_path" --light --test-v2 --upgrade-gateways ${COMMON_ARGS}
+ zetae2e local $E2E_ARGS --skip-setup --config "$deployed_config_path" --light --test-v2 --upgrade-contracts ${COMMON_ARGS}
  else
- zetae2e local $E2E_ARGS --skip-setup --config "$deployed_config_path" --skip-bitcoin-setup --light --test-v2 --upgrade-gateways ${COMMON_ARGS}
+ zetae2e local $E2E_ARGS --skip-setup --config "$deployed_config_path" --skip-bitcoin-setup --light --test-v2 --upgrade-contracts ${COMMON_ARGS}
  fi
 
  ZETAE2E_EXIT_CODE=$?

e2e/runner/v2_gateway.go → e2e/runner/v2_upgrade.go

@@ -3,17 +3,19 @@ package runner
 import (
  ethtypes "github.com/ethereum/go-ethereum/core/types"
  "github.com/stretchr/testify/require"
+ "github.com/zeta-chain/protocol-contracts/v2/pkg/erc20custody.sol"
  "github.com/zeta-chain/protocol-contracts/v2/pkg/gatewayevm.sol"
  "github.com/zeta-chain/protocol-contracts/v2/pkg/gatewayzevm.sol"
 
  "github.com/zeta-chain/node/e2e/utils"
 )
 
-// UpgradeGateways upgrades the GatewayEVM and GatewayZEVM contracts
-// It deploy new gateway contract implementation with the current imported artifacts and upgrade the gateway contract
-func (r *E2ERunner) UpgradeGateways() {
+// UpgradeContracts upgrades contracts
+// It deploys new contract implementation with the current imported artifacts and upgrade the contract
+func (r *E2ERunner) UpgradeContracts() {
  r.UpgradeGatewayZEVM()
  r.UpgradeGatewayEVM()
+ r.UpgradeERC20Custody()
 }
 
 // UpgradeGatewayZEVM upgrades the GatewayZEVM contract
@@ -53,3 +55,22 @@ func (r *E2ERunner) UpgradeGatewayEVM() {
  require.NoError(r, err)
  ensureTxReceipt(txUpgrade, "GatewayEVM upgrade failed")
 }
+
+// UpgradeERC20CustodyZEVM upgrades the ERC20Custody contract
+func (r *E2ERunner) UpgradeERC20Custody() {
+ ensureTxReceipt := func(tx *ethtypes.Transaction, failMessage string) {
+ receipt := utils.MustWaitForTxReceipt(r.Ctx, r.EVMClient, tx, r.Logger, r.ReceiptTimeout)
+ r.requireTxSuccessful(receipt, failMessage+" tx hash: "+tx.Hash().Hex())
+ }
+
+ r.Logger.Info("Upgrading ERC20Custody contract")
+ // Deploy the new erc20Custody contract implementation
+ newImplementationAddress, txDeploy, _, err := erc20custody.DeployERC20Custody(r.EVMAuth, r.EVMClient)
+ require.NoError(r, err)
+ ensureTxReceipt(txDeploy, "New ERC20Custody implementation deployment failed")
+
+ // Upgrade
+ txUpgrade, err := r.ERC20CustodyV2.UpgradeToAndCall(r.EVMAuth, newImplementationAddress, []byte{})
+ require.NoError(r, err)
+ ensureTxReceipt(txUpgrade, "ERC20Custody upgrade failed")
+}