refactor(test:zk): simplify getCode

matter-labs · Oct 15, 2024 · bf238ab · bf238ab
1 parent 8d70814
commit bf238ab
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 25 deletions.
diff --git a/crates/forge/tests/it/zk/cheats.rs b/crates/forge/tests/it/zk/cheats.rs
@@ -18,21 +18,10 @@ async fn test_zk_cheat_get_code() {
     let mut zk_config = TEST_DATA_DEFAULT.zk_test_data.zk_config.clone();
     zk_config
         .fs_permissions
-        .add(PathPermission::read_write("./zk/zkout/ConstantNumber.sol/ConstantNumber.json"));
-    let runner = TEST_DATA_DEFAULT.runner_with_zksync_config(zk_config);
-    let filter = Filter::new("testZkCheatcodesGetCode", "ZkCheatcodesTest", ".*");
-
-    TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
-}
+        .add(PathPermission::read("./zk"));
 
-#[tokio::test(flavor = "multi_thread")]
-async fn test_zk_cheat_get_code_evm() {
-    let mut zk_config = TEST_DATA_DEFAULT.zk_test_data.zk_config.clone();
-    zk_config
-        .fs_permissions
-        .add(PathPermission::read_write("./zk/out/ConstantNumber.sol/ConstantNumber.json"));
     let runner = TEST_DATA_DEFAULT.runner_with_zksync_config(zk_config);
-    let filter = Filter::new("testZkCheatcodesGetEVMCode", "ZkCheatcodesTest", ".*");
+    let filter = Filter::new("testZkCheatcodesGetCode", "ZkCheatcodesTest", ".*");
 
     TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
 }

diff --git a/testdata/zk/Cheatcodes.t.sol b/testdata/zk/Cheatcodes.t.sol
@@ -127,24 +127,22 @@ contract ZkCheatcodesTest is DSTest {
         require(vm.getNonce(TEST_ADDRESS) == 0, "era nonce mismatch");
     }
 
-    function testZkCheatcodesGetCode() public {
-        bytes memory fullPath = vm.getCode("ConstantNumber.sol");
+    function getCodeCheck(string memory contractName, string memory outDir) internal {
+        bytes memory bytecode = vm.getCode(contractName);
 
-        string memory artifact = vm.readFile("zk/zkout/ConstantNumber.sol/ConstantNumber.json");
+        string memory artifactPath = string.concat("zk/", outDir, "/", contractName, ".sol/", contractName, ".json");
+        string memory artifact = vm.readFile(artifactPath);
         bytes memory expectedBytecode = vm.parseJsonBytes(artifact, ".bytecode.object");
 
-        assertEq(fullPath, expectedBytecode, "code for the contract was incorrect");
+        assertEq(bytecode, expectedBytecode, "code for the contract was incorrect");
     }
 
-    function testZkCheatcodesGetEVMCode() public {
-        vm.zkVm(false);
-        bytes memory fullPath = vm.getCode("ConstantNumber.sol");
-        vm.zkVm(true);
-
-        string memory artifact = vm.readFile("zk/out/ConstantNumber.sol/ConstantNumber.json");
-        bytes memory expectedBytecode = vm.parseJsonBytes(artifact, ".bytecode.object");
+    function testZkCheatcodesGetCode() public {
+        string memory contractName = "ConstantNumber";
+        getCodeCheck(contractName, "zkout");
 
-        assertEq(fullPath, expectedBytecode, "code for the contract was incorrect");
+        vm.zkVm(false);
+        getCodeCheck(contractName, "out");
     }
 
     function testZkCheatcodesEtch() public {