Skip to content

Commit d83b606

Browse files
author
Ryan
committed
feat: add unit tests
1 parent 584ee46 commit d83b606

5 files changed

Lines changed: 853 additions & 0 deletions

File tree

‎test/ERC1967/beacon/BeaconProxy.t.sol‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity ^0.8.24;
3+
4+
import {Test} from "forge-std/Test.sol";
5+
6+
import {ERC1967Utils} from "src/ERC1967/ERC1967Utils.sol";
7+
import {BeaconProxy} from "src/ERC1967/beacon/BeaconProxy.sol";
8+
import {UpgradeableBeacon} from "src/ERC1967/beacon/UpgradeableBeacon.sol";
9+
10+
import {MockImplementationV1, MockImplementationV2} from "test/mocks/MockImplementation.sol";
11+
12+
contract BeaconProxyTest is Test {
13+
error Unauthorized();
14+
15+
address internal immutable alice = makeAddr("alice");
16+
address internal immutable bob = makeAddr("bob");
17+
18+
UpgradeableBeacon internal beacon;
19+
20+
address payable internal proxy;
21+
22+
MockImplementationV1 internal mockV1;
23+
MockImplementationV2 internal mockV2;
24+
25+
MockImplementationV1 internal implementationV1;
26+
MockImplementationV2 internal implementationV2;
27+
28+
function setUp() public virtual {
29+
implementationV1 = new MockImplementationV1();
30+
implementationV2 = new MockImplementationV2();
31+
32+
beacon = new UpgradeableBeacon(address(implementationV1), address(this));
33+
34+
bytes memory data = abi.encodeWithSelector(MockImplementationV1.initialize.selector, uint256(10));
35+
proxy = payable(address(new BeaconProxy(address(beacon), data)));
36+
37+
mockV1 = MockImplementationV1(proxy);
38+
mockV2 = MockImplementationV2(proxy);
39+
}
40+
41+
// ============================================================
42+
// Constructor
43+
// ============================================================
44+
45+
function test_constructor_setsBeaconSlot() public view {
46+
assertEq(_getBeacon(proxy), address(beacon));
47+
}
48+
49+
function test_constructor_executesInitializer() public view {
50+
assertEq(mockV1.getNumber(), uint256(10));
51+
assertEq(mockV1.owner(), address(this));
52+
assertEq(mockV1.getVersion(), uint64(1));
53+
}
54+
55+
function test_constructor_resolvesImplementationViaBeacon() public view {
56+
assertEq(beacon.implementation(), address(implementationV1));
57+
}
58+
59+
function test_constructor_emitsBeaconUpgraded() public {
60+
bytes memory data = abi.encodeWithSelector(MockImplementationV1.initialize.selector, uint256(1));
61+
62+
vm.expectEmit(true, false, false, false);
63+
emit ERC1967Utils.BeaconUpgraded(address(beacon));
64+
new BeaconProxy(address(beacon), data);
65+
}
66+
67+
function test_constructor_revertsIfInvalidBeacon() public {
68+
bytes memory data = abi.encodeWithSelector(MockImplementationV1.initialize.selector, uint256(1));
69+
70+
vm.expectRevert(ERC1967Utils.InvalidBeacon.selector);
71+
new BeaconProxy(address(0xdead), data);
72+
}
73+
74+
function test_constructor_revertsIfBeaconReturnsEOA() public {
75+
// Deploy a beacon pointing to an EOA — the beacon itself deploys, but
76+
// creating a BeaconProxy against it should fail because the implementation
77+
// returned by the beacon has no code.
78+
vm.expectRevert();
79+
new UpgradeableBeacon(alice, address(this));
80+
}
81+
82+
// ============================================================
83+
// Delegation
84+
// ============================================================
85+
86+
function test_delegation_delegatesToImplementation() public view {
87+
assertEq(mockV1.getNumber(), uint256(10));
88+
}
89+
90+
function test_delegation_stateIsPersisted() public {
91+
mockV1.setNumber(42);
92+
assertEq(mockV1.getNumber(), uint256(42));
93+
}
94+
95+
function test_delegation_increment() public {
96+
mockV1.increment();
97+
assertEq(mockV1.getNumber(), uint256(11));
98+
}
99+
100+
function test_delegation_revertsIfUnauthorized() public {
101+
vm.prank(alice);
102+
vm.expectRevert();
103+
mockV1.setNumber(99);
104+
}
105+
106+
function test_delegation_acceptsEther() public {
107+
deal(alice, 1 ether);
108+
vm.prank(alice);
109+
(bool success,) = proxy.call{value: 0.5 ether}("");
110+
assertTrue(success);
111+
assertEq(proxy.balance, 0.5 ether);
112+
}
113+
114+
// ============================================================
115+
// Beacon upgrade
116+
// ============================================================
117+
118+
function test_upgradeTo_updatesImplementation() public {
119+
beacon.upgradeTo(address(implementationV2));
120+
assertEq(beacon.implementation(), address(implementationV2));
121+
}
122+
123+
function test_upgradeTo_proxyDelegatesNewImplementation() public {
124+
mockV1.setNumber(20);
125+
126+
bytes memory data = abi.encodeWithSelector(MockImplementationV2.initialize.selector, uint256(5));
127+
beacon.upgradeTo(address(implementationV2));
128+
129+
// Reinitialize V2 through the proxy
130+
(bool success,) = proxy.call(data);
131+
assertTrue(success);
132+
133+
assertEq(mockV2.getNumber(), uint256(20));
134+
assertEq(mockV2.getMultiplier(), uint256(5));
135+
}
136+
137+
function test_upgradeTo_newBehaviorAfterUpgrade() public {
138+
mockV1.setNumber(10);
139+
140+
beacon.upgradeTo(address(implementationV2));
141+
142+
bytes memory data = abi.encodeWithSelector(MockImplementationV2.initialize.selector, uint256(3));
143+
(bool success,) = proxy.call(data);
144+
assertTrue(success);
145+
146+
mockV2.increment();
147+
assertEq(mockV2.getNumber(), uint256(13));
148+
}
149+
150+
function test_upgradeTo_preservesState() public {
151+
mockV1.setNumber(42);
152+
mockV1.increment();
153+
assertEq(mockV1.getNumber(), uint256(43));
154+
155+
beacon.upgradeTo(address(implementationV2));
156+
assertEq(mockV2.getNumber(), uint256(43));
157+
}
158+
159+
function test_upgradeTo_emitsUpgraded() public {
160+
vm.expectEmit(true, false, false, false, address(beacon));
161+
emit UpgradeableBeacon.Upgraded(address(implementationV2));
162+
beacon.upgradeTo(address(implementationV2));
163+
}
164+
165+
function test_upgradeTo_revertsIfNotOwner() public {
166+
vm.prank(alice);
167+
vm.expectRevert(Unauthorized.selector);
168+
beacon.upgradeTo(address(implementationV2));
169+
}
170+
171+
function test_upgradeTo_revertsIfInvalidImplementation() public {
172+
vm.expectRevert(UpgradeableBeacon.InvalidBeaconImplementation.selector);
173+
beacon.upgradeTo(address(0xdead));
174+
}
175+
176+
function test_upgradeTo_revertsIfEOAImplementation() public {
177+
vm.expectRevert(UpgradeableBeacon.InvalidBeaconImplementation.selector);
178+
beacon.upgradeTo(alice);
179+
}
180+
181+
// ============================================================
182+
// Multiple proxies sharing one beacon
183+
// ============================================================
184+
185+
function test_multipleProxies_shareBeaconImplementation() public {
186+
bytes memory data1 = abi.encodeWithSelector(MockImplementationV1.initialize.selector, uint256(100));
187+
bytes memory data2 = abi.encodeWithSelector(MockImplementationV1.initialize.selector, uint256(200));
188+
189+
address payable proxy1 = payable(address(new BeaconProxy(address(beacon), data1)));
190+
address payable proxy2 = payable(address(new BeaconProxy(address(beacon), data2)));
191+
192+
MockImplementationV1 mock1 = MockImplementationV1(proxy1);
193+
MockImplementationV1 mock2 = MockImplementationV1(proxy2);
194+
195+
assertEq(mock1.getNumber(), uint256(100));
196+
assertEq(mock2.getNumber(), uint256(200));
197+
198+
// Upgrade beacon — both proxies get the new implementation.
199+
beacon.upgradeTo(address(implementationV2));
200+
201+
MockImplementationV2 mockV2_1 = MockImplementationV2(proxy1);
202+
MockImplementationV2 mockV2_2 = MockImplementationV2(proxy2);
203+
204+
// State is preserved independently.
205+
assertEq(mockV2_1.getNumber(), uint256(100));
206+
assertEq(mockV2_2.getNumber(), uint256(200));
207+
}
208+
209+
// ============================================================
210+
// Beacon constructor with empty init data
211+
// ============================================================
212+
213+
function test_constructor_withEmptyData() public {
214+
// BeaconProxy does not have the ProxyUninitialized check (that's on ERC1967Proxy).
215+
// Empty data should succeed as long as no value is sent.
216+
address payable p = payable(address(new BeaconProxy(address(beacon), "")));
217+
assertTrue(p.code.length > 0);
218+
}
219+
220+
function test_constructor_revertsNonPayableIfValueWithEmptyData() public {
221+
vm.deal(address(this), 1 ether);
222+
vm.expectRevert(ERC1967Utils.NonPayable.selector);
223+
new BeaconProxy{value: 1 wei}(address(beacon), "");
224+
}
225+
226+
// ============================================================
227+
// Helpers
228+
// ============================================================
229+
230+
function _getBeacon(address target) internal view returns (address) {
231+
return address(uint160(uint256(vm.load(target, ERC1967Utils.BEACON_SLOT))));
232+
}
233+
}

0 commit comments

Comments
 (0)