fix: Add handle expect emit on create (#402)

Co-authored-by: Nisheeth Barthwal <[email protected]> Co-authored-by: Herman Obst Demaestri <[email protected]>
matter-labs · Jun 3, 2024 · d94bab2 · d94bab2
1 parent cd95784
commit d94bab2
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 1 deletion.
diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs
@@ -1738,13 +1738,20 @@ impl<DB: DatabaseExt + Send> Inspector<DB> for Cheatcodes {
                 &mut data.journaled_state,
                 ccx,
             ) {
-                self.combined_logs.extend(result.logs.into_iter().map(|log| {
+                self.combined_logs.extend(result.logs.clone().into_iter().map(|log| {
                     Some(Log {
                         address: log.address,
                         data: LogData::new_unchecked(log.topics, log.data),
                     })
                 }));
 
+                // for each log in cloned logs call handle_expect_emit
+                if !self.expected_emits.is_empty() {
+                    for log in result.logs {
+                        expect::handle_expect_emit(self, &log.address, &log.topics, &log.data);
+                    }
+                }
+
                 return match result.execution_result {
                     ExecutionResult::Success { output, .. } => match output {
                         Output::Create(bytes, address) => {

diff --git a/zk-tests/src/Cheatcodes.t.sol b/zk-tests/src/Cheatcodes.t.sol
@@ -36,6 +36,7 @@ interface IMyProxyCaller {
 
 contract MyProxyCaller {
     address inner;
+
     constructor(address _inner) {
         inner = _inner;
     }
@@ -45,7 +46,22 @@ contract MyProxyCaller {
     }
 }
 
+contract Emitter {
+    event EventConstructor(string message);
+    event EventFunction(string message);
+
+    constructor() {
+        emit EventConstructor("constructor");
+    }
+
+    function functionEmit() public {
+        emit EventFunction("function");
+    }
+}
+
 contract ZkCheatcodesTest is Test {
+    event EventConstructor(string message);
+    event EventFunction(string message);
     uint256 testSlot = 0; //0x000000000000000000000000000000000000000000000000000000000000001e slot
     uint256 constant ERA_FORK_BLOCK = 19579636;
     uint256 constant ERA_FORK_BLOCK_TS = 1700601590;
@@ -143,6 +159,19 @@ contract ZkCheatcodesTest is Test {
         assertEq(writes[0], keySlot0);
     }
 
+    function testExpectEmit() public {
+        vm.expectEmit(true, true, true, true);
+        emit EventFunction("function");
+        Emitter emitter = new Emitter();
+        emitter.functionEmit();
+    }
+
+    function testExpectEmitOnCreate() public {
+        vm.expectEmit(true, true, true, true);
+        emit EventConstructor("constructor");
+        new Emitter();
+    }
+
     function testZkCheatcodesValueFunctionMockReturn() public {
         InnerMock inner = new InnerMock();
         // Send some funds to so it can pay for the inner call