使用 Wake 框架安全地测试代理合约

介绍

升级是生产 bug 藏身之处：遗漏的初始化程序、错误的管理员或损坏的存储。代理模式允许你升级合约，但它们引入了传统测试所遗漏的复杂性。Wake 基于 Python 的测试可以在这些问题到达主网之前发现它们。

结果是干净的测试代码。通过代理调用实现函数非常简单：

contract = ExampleERC20Upgradeable(proxy)

以下是如何在 Wake 中测试代理合约。

1. 导入代理合约

Wake 需要编译你的代理合约以生成 Python 类型绑定（pytypes）。如果代理合约位于你的库目录中，Wake 默认不会编译它。

在 wake.toml 中，默认配置是 exclude_paths = ["script", ".venv", "venv", "node_modules", "lib", "test"]。这意味着除非从非排除的文件导入，否则不会编译这些路径中的合约。

要使合约在你的项目中可用，请从 exclude_paths 之外的某个位置导入它。有关更多详细信息，请参见文档：https://ackee.xyz/wake/docs/latest/compilation/

创建 tests/imports.sol 以使 pytypes 可用：

import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";

2. 在 Python 中导入代理

再次运行 wake up 以编译。Wake 为你的实现和代理合约生成 Python 绑定。将它们导入到你的测试文件中：

tests/test_upgradable.py

from pytypes.contracts.ExampleERC20Upgradeable import ExampleERC20Upgradeable
from pytypes.openzeppelin.contracts.proxy.ERC1967.ERC1967Proxy import ERC1967Proxy

3. 部署和初始化

首先部署实现合约，然后创建指向它的代理。代理的初始化数据编码了对实现的 initialize 函数的调用：

@chain.connect()
@on_revert(revert_handler)
def test_default():

    impl_erc20 = ExampleERC20Upgradeable.deploy()

    proxy = ERC1967Proxy.deploy(
        implementation =impl_erc20,
        _data=abi.encode_call(ExampleERC20Upgradeable.initialize, ("Upgradable Token", "UPG", 10**20, chain.accounts[0])),
        from_=chain.accounts[0]
    )

_data 参数编码了在代理部署期间运行的初始化调用。这取代了非可升级合约中使用的构造函数模式。

4. 通过代理访问实现函数

用实现合约类包装代理地址。这告诉 Wake 通过代理路由所有函数调用，同时使用实现的 ABI：

contract = ExampleERC20Upgradeable(proxy)

Wake 自动处理 delegatecall 路由，因此你可以像与简单部署一样与合约交互。

5. 调用实现函数

所有实现函数现在都通过包装的代理可用。你可以验证合约的行为，检查事件并测试状态更改：

## 验证初始余额
assert contract.balanceOf(chain.accounts[1]) == 0

## 执行转账
tx = contract.transfer(chain.accounts[1], 10**18, from_=chain.accounts[0])

## 检查发出的事件
event = next(event for event in tx.events if isinstance(event, ExampleERC20Upgradeable.Transfer))
assert event.from_ == chain.accounts[0].address
assert event.to == chain.accounts[1].address
assert event.value == 10**18

## 验证更新后的余额
assert contract.balanceOf(chain.accounts[1]) == 10**18

该测试确认代理正确地委托给实现并按预期维护状态。

结论

Wake 通过其 Python 绑定简化了代理测试。用实现类包装代理地址并直接调用函数。相同的方法适用于单元测试和手动引导模糊测试 (MGF)，使你可以使用与标准合约相同的工具来测试可升级合约。

这可以在升级 bug（遗漏的初始化程序、存储冲突、访问控制问题）变成漏洞之前捕获它们。以与测试其他所有内容相同的方式测试你的代理模式。

在此处了解更多信息，请查看：手动引导模糊测试初学者指南

附录：完整测试代码

import math
from wake.testing import *
from dataclasses import dataclass

from pytypes.contracts.ExampleERC20Upgradeable import ExampleERC20Upgradeable
from pytypes.openzeppelin.contracts.proxy.ERC1967.ERC1967Proxy import ERC1967Proxy

## 打印失败的 tx 调用跟踪
def revert_handler(e: RevertError):
    if e.tx is not None:
        print(e.tx.call_trace)

@chain.connect()
@on_revert(revert_handler)
def test_default():

    impl_erc20 = ExampleERC20Upgradeable.deploy()

    proxy = ERC1967Proxy.deploy(
        implementation =impl_erc20,
        _data=abi.encode_call(ExampleERC20Upgradeable.initialize, ("Upgradable Token", "UPG", 10**20, chain.accounts[0])),
        from_=chain.accounts[0]
    )

    contract = ExampleERC20Upgradeable(proxy) # Just wrap the proxy with the contract Class to call functions

    assert contract.balanceOf(chain.accounts[1]) == 0
    tx = contract.transfer(chain.accounts[1], 10**18, from_=chain.accounts[0])

    event = next(event for event in tx.events if isinstance(event, ExampleERC20Upgradeable.Transfer))
    assert event.from_ == chain.accounts[0].address
    assert event.to == chain.accounts[1].address
    assert event.value == 10**18

    assert contract.balanceOf(chain.accounts[1]) == 10**18

• 原文链接： ackee.xyz/blog/test-prox...
• 登链社区 AI 助手，为大家转译优秀英文文章，如有翻译不通的地方，还请包涵～