Merge pull request #8 from base-org/amie/test-upgrades

Test `upgradeToAndCall` for ERC-1967 upgrades

Author: amiecorso

.gitmodules

@@ -7,3 +7,6 @@
 [submodule "lib/smart-wallet"]
 	path = lib/smart-wallet
 	url = https://github.com/coinbase/smart-wallet
+[submodule "lib/solady"]
+	path = lib/solady
+	url = https://github.com/vectorized/solady

lib/solady

@@ -1,3 +1,4 @@
 smart-wallet/=lib/smart-wallet/src/
 forge-std/=lib/forge-std/src/
-openzeppelin-contracts/=lib/openzeppelin-contracts/ 
+openzeppelin-contracts/=lib/openzeppelin-contracts/
+solady/=lib/solady/src/

remappings.txt

@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.23;
+
+import {EIP7702ProxyBase} from "../base/EIP7702ProxyBase.sol";
+import {EIP7702Proxy} from "../../src/EIP7702Proxy.sol";
+import {CoinbaseSmartWallet} from "../../lib/smart-wallet/src/CoinbaseSmartWallet.sol";
+import {UUPSUpgradeable} from "solady/utils/UUPSUpgradeable.sol";
+
+contract UpgradeToAndCallTest is EIP7702ProxyBase {
+    DummyImplementation newImplementation;
+
+    function setUp() public override {
+        super.setUp();
+
+        // Initialize the proxy first
+        bytes memory initArgs = _createInitArgs(_newOwner);
+        bytes memory signature = _signInitData(_EOA_PRIVATE_KEY, initArgs);
+        EIP7702Proxy(_eoa).initialize(initArgs, signature);
+
+        // Deploy new implementation
+        newImplementation = new DummyImplementation();
+    }
+
+    function testUpgradeToAndCall_succeedsForOwner() public {
+        // Wallet owner should be able to upgrade
+        vm.prank(_newOwner);
+
+        CoinbaseSmartWallet(payable(_eoa)).upgradeToAndCall(
+            address(newImplementation),
+            abi.encodeWithSignature("dummy()")
+        );
+
+        // Verify upgrade worked by calling new function
+        vm.expectEmit(true, true, true, true, _eoa);
+        emit DummyImplementation.DummyCalled();
+        DummyImplementation(payable(_eoa)).dummy();
+    }
+
+    function testUpgradeToAndCall_succeedsForEOA() public {
+        // EOA should be able to upgrade
+        vm.prank(_eoa);
+
+        CoinbaseSmartWallet(payable(_eoa)).upgradeToAndCall(
+            address(newImplementation),
+            abi.encodeWithSignature("dummy()")
+        );
+
+        // Verify upgrade worked by calling new function
+        vm.expectEmit(true, true, true, true, _eoa);
+        emit DummyImplementation.DummyCalled();
+        DummyImplementation(payable(_eoa)).dummy();
+    }
+
+    function testUpgradeToAndCall_revertsForNonOwner() public {
+        vm.prank(address(0xBAD));
+        vm.expectRevert(); // CoinbaseSmartWallet will revert for non-owner
+        CoinbaseSmartWallet(payable(_eoa)).upgradeToAndCall(
+            address(newImplementation),
+            ""
+        );
+    }
+}
+
+contract DummyImplementation is UUPSUpgradeable {
+    event DummyCalled();
+
+    function dummy() external {
+        emit DummyCalled();
+    }
+
+    function _authorizeUpgrade(address) internal override {}
+}