feat(cheats): deployCodeTo (#401)

* feat(cheats): `deployCodeTo` * Update src/StdCheats.sol Co-authored-by: Matt Solomon <[email protected]> * test: fix test --------- Co-authored-by: Matt Solomon <[email protected]>
foundry-rs · Jun 25, 2023 · 2f43c7e · 2f43c7e
1 parent 27035c8
commit 2f43c7e
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 0 deletions.
diff --git a/src/StdCheats.sol b/src/StdCheats.sol
@@ -697,4 +697,20 @@ abstract contract StdCheats is StdCheatsSafe {
         // update owner
         stdstore.target(token).sig(0x6352211e).with_key(id).checked_write(to);
     }
+
+    function deployCodeTo(string memory what, address where) internal virtual {
+        deployCodeTo(what, "", 0, where);
+    }
+
+    function deployCodeTo(string memory what, bytes memory args, address where) internal virtual {
+        deployCodeTo(what, args, 0, where);
+    }
+
+    function deployCodeTo(string memory what, bytes memory args, uint256 value, address where) internal virtual {
+        bytes memory creationCode = vm.getCode(what);
+        vm.etch(where, abi.encodePacked(creationCode, args));
+        (bool success, bytes memory runtimeBytecode) = where.call{value: value}("");
+        require(success, "StdCheats deployCodeTo(string,bytes,uint256,address): Failed to create runtime bytecode.");
+        vm.etch(where, runtimeBytecode);
+    }
 }
diff --git a/test/StdCheats.t.sol b/test/StdCheats.t.sol
@@ -366,6 +366,33 @@ contract StdCheatsTest is Test {
                 && addr != 0x4e59b44847b379578588920cA78FbF26c0B4956C
         );
     }
+
+    function testCannotDeployCodeTo() external {
+        vm.expectRevert("StdCheats deployCodeTo(string,bytes,uint256,address): Failed to create runtime bytecode.");
+        this._revertDeployCodeTo();
+    }
+
+    function _revertDeployCodeTo() external {
+        deployCodeTo("StdCheats.t.sol:RevertingContract", address(0));
+    }
+
+    function testDeployCodeTo() external {
+        address arbitraryAddress = makeAddr("arbitraryAddress");
+
+        deployCodeTo(
+            "StdCheats.t.sol:MockContractWithConstructorArgs",
+            abi.encode(uint256(6), true, bytes20(arbitraryAddress)),
+            1 ether,
+            arbitraryAddress
+        );
+
+        MockContractWithConstructorArgs ct = MockContractWithConstructorArgs(arbitraryAddress);
+
+        assertEq(arbitraryAddress.balance, 1 ether);
+        assertEq(ct.x(), 6);
+        assertTrue(ct.y());
+        assertEq(ct.z(), bytes20(arbitraryAddress));
+    }
 }
 
 contract StdCheatsMock is StdCheats {
@@ -505,3 +532,15 @@ contract RevertingContract {
         revert();
     }
 }
+
+contract MockContractWithConstructorArgs {
+    uint256 public immutable x;
+    bool public y;
+    bytes20 public z;
+
+    constructor(uint256 _x, bool _y, bytes20 _z) payable {
+        x = _x;
+        y = _y;
+        z = _z;
+    }
+}