fix: test cheatcodes in contract

Moonsong-Labs · Nov 14, 2023 · 547f362 · 547f362
1 parent c1690ae
commit 547f362
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/e2e-tests/contracts/TestCheatcodes.sol b/e2e-tests/contracts/TestCheatcodes.sol
@@ -5,14 +5,20 @@ pragma solidity ^0.8.0;
 contract TestCheatcodes {
   address constant CHEATCODE_ADDRESS = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
 
-  function deal(address account, uint256 amount) external {
+  function testDeal(address account, uint256 amount) external {
+    uint balanceBefore = address(account).balance;
     (bool success, ) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("deal(address,uint256)", account, amount));
+    uint balanceAfter = address(account).balance;
+    require(balanceAfter == amount, "balance mismatch");
+    require(balanceAfter != balanceBefore, "balance unchanged");
     require(success, "deal failed");
   }
 
-  function etch(address target, bytes calldata code) external {
+  function testEtch(address target, bytes calldata code) external {
     (bool success, ) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("etch(address,bytes)", target, code));
     require(success, "etch failed");
+    (success, ) = target.call(abi.encodeWithSignature("setGreeting(bytes)", bytes("hello world")));
+    require(success, "setGreeting failed");
   }
 
   function setNonce(address account, uint256 nonce) external {

diff --git a/e2e-tests/test/cheatcodes.test.ts b/e2e-tests/test/cheatcodes.test.ts
@@ -13,18 +13,14 @@ describe("Cheatcodes", function () {
     const wallet = new Wallet(RichAccounts[0].PrivateKey);
     const deployer = new Deployer(hre, wallet);
     const randomWallet = Wallet.createRandom().connect(provider);
-    const initialBalance = await provider.getBalance(randomWallet.address);
 
     // Act
     const greeter = await deployContract(deployer, "TestCheatcodes", []);
-    await greeter.deal(randomWallet.address, 123456, {
+    await greeter.testDeal(randomWallet.address, 123456, {
       gasLimit: 1000000,
     });
 
-    // Assert
-    const finalBalance = await provider.getBalance(randomWallet.address);
-    expect(finalBalance.toNumber()).to.eq(123456);
-    expect(finalBalance).to.not.eq(initialBalance);
+    // Asserts are in contract
   });
 
   it("Should test vm.etch", async function () {
@@ -38,7 +34,7 @@ describe("Cheatcodes", function () {
     const cheatcodes = await deployContract(deployer, "TestCheatcodes", []);
     const greeter = await deployContract(deployer, "Greeter", ["Hi"]);
     const greeterCode = await provider.getCode(greeter.address);
-    await cheatcodes.etch(randomWallet.address, greeterCode);
+    await cheatcodes.testEtch(randomWallet.address, greeterCode);
 
     // Assert
     expect(initialRandomWalletCode).to.eq("0x");