From bf238ab05311e6a537ab9bffdc0f58993ad11286 Mon Sep 17 00:00:00 2001
From: Francesco Dainese <franci.dainese@gmail.com>
Date: Tue, 15 Oct 2024 17:23:10 -0300
Subject: [PATCH] refactor(test:zk): simplify getCode

---
 crates/forge/tests/it/zk/cheats.rs | 15 ++-------------
 testdata/zk/Cheatcodes.t.sol       | 22 ++++++++++------------
 2 files changed, 12 insertions(+), 25 deletions(-)

diff --git a/crates/forge/tests/it/zk/cheats.rs b/crates/forge/tests/it/zk/cheats.rs
index f651deccd..7b6110123 100644
--- 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
index 1acabf94e..ba4c1394a 100644
--- 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 {