feat(cheats): json/toml --> consolidate write and writeUpsert (#11447)

Author: 0xrusowsky

crates/cheatcodes/assets/cheatcodes.json

@@ -2552,15 +2552,10 @@ interface Vm {
 
     /// Write a serialized JSON object to an **existing** JSON file, replacing a value with key = <value_key.>
     /// This is useful to replace a specific value of a JSON file, without having to parse the entire thing.
+    /// This cheatcode will create new keys if they didn't previously exist.
     #[cheatcode(group = Json)]
     function writeJson(string calldata json, string calldata path, string calldata valueKey) external;
 
-    /// Write a serialized JSON object to an **existing** JSON file, replacing a value with key = <value_key.>
-    /// This is useful to replace a specific value of a JSON file, without having to parse the entire thing.
-    /// Unlike `writeJson`, this cheatcode will create new keys if they didn't previously exist.
-    #[cheatcode(group = Toml)]
-    function writeJsonUpsert(string calldata json, string calldata path, string calldata valueKey) external;
-
     // ======== TOML Parsing and Manipulation ========
 
     // -------- Reading --------
@@ -2663,15 +2658,10 @@ interface Vm {
 
     /// Takes serialized JSON, converts to TOML and write a serialized TOML table to an **existing** TOML file, replacing a value with key = <value_key.>
     /// This is useful to replace a specific value of a TOML file, without having to parse the entire thing.
+    /// This cheatcode will create new keys if they didn't previously exist.
     #[cheatcode(group = Toml)]
     function writeToml(string calldata json, string calldata path, string calldata valueKey) external;
 
-    /// Takes serialized JSON, converts to TOML and write a serialized TOML table to an **existing** TOML file, replacing a value with key = <value_key.>
-    /// This is useful to replace a specific value of a TOML file, without having to parse the entire thing.
-    /// Unlike `writeToml`, this cheatcode will create new keys if they didn't previously exist.
-    #[cheatcode(group = Toml)]
-    function writeTomlUpsert(string calldata json, string calldata path, string calldata valueKey) external;
-
     // ======== Cryptography ========
 
     // -------- Key Management --------

crates/cheatcodes/spec/src/vm.rs

@@ -344,24 +344,6 @@ impl Cheatcode for writeJson_0Call {
 }
 
 impl Cheatcode for writeJson_1Call {
-    fn apply(&self, state: &mut Cheatcodes) -> Result {
-        let Self { json, path, valueKey } = self;
-        let json = serde_json::from_str(json).unwrap_or_else(|_| Value::String(json.to_owned()));
-
-        let data_path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
-        let data_s = fs::read_to_string(data_path)?;
-        let data = serde_json::from_str(&data_s)?;
-        let value =
-            jsonpath_lib::replace_with(data, &canonicalize_json_path(valueKey), &mut |_| {
-                Some(json.clone())
-            })?;
-
-        let json_string = serde_json::to_string_pretty(&value)?;
-        super::fs::write_file(state, path.as_ref(), json_string.as_bytes())
-    }
-}
-
-impl Cheatcode for writeJsonUpsertCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { json: value, path, valueKey } = self;
 

crates/cheatcodes/src/json.rs

@@ -4,8 +4,8 @@ use crate::{
     Cheatcode, Cheatcodes, Result,
     Vm::*,
     json::{
-        canonicalize_json_path, check_json_key_exists, parse_json, parse_json_coerce,
-        parse_json_keys, resolve_type, upsert_json_value,
+        check_json_key_exists, parse_json, parse_json_coerce, parse_json_keys, resolve_type,
+        upsert_json_value,
     },
 };
 use alloy_dyn_abi::DynSolType;
@@ -175,26 +175,6 @@ impl Cheatcode for writeToml_0Call {
 }
 
 impl Cheatcode for writeToml_1Call {
-    fn apply(&self, state: &mut Cheatcodes) -> Result {
-        let Self { json, path, valueKey } = self;
-        let json =
-            serde_json::from_str(json).unwrap_or_else(|_| JsonValue::String(json.to_owned()));
-
-        let data_path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
-        let toml_data = fs::read_to_string(data_path)?;
-        let json_data: JsonValue =
-            toml::from_str(&toml_data).map_err(|e| fmt_err!("failed parsing TOML: {e}"))?;
-        let value =
-            jsonpath_lib::replace_with(json_data, &canonicalize_json_path(valueKey), &mut |_| {
-                Some(json.clone())
-            })?;
-
-        let toml_string = format_json_to_toml(value)?;
-        super::fs::write_file(state, path.as_ref(), toml_string.as_bytes())
-    }
-}
-
-impl Cheatcode for writeTomlUpsertCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { json: value, path, valueKey } = self;
 

crates/cheatcodes/src/toml.rs

@@ -491,4 +491,28 @@ contract WriteJsonTest is DSTest {
         address decodedAddress = abi.decode(data, (address));
         assertEq(decodedAddress, ex);
     }
+
+    function test_writeJson_createKeys() public {
+        string memory path = "fixtures/Json/write_test.json";
+        string memory json = vm.readFile(path);
+
+        bool exists = vm.keyExistsJson(json, ".parent");
+        assertTrue(!exists);
+        exists = vm.keyExistsJson(json, ".parent.child");
+        assertTrue(!exists);
+        exists = vm.keyExistsJson(json, ".parent.child.value");
+        assertTrue(!exists);
+
+        // Write to nested path, creating intermediate keys
+        vm.writeJson(vm.toString(uint256(42)), path, ".parent.child.value");
+
+        // Verify the value was written and intermediate keys were created
+        json = vm.readFile(path);
+        uint256 value = abi.decode(vm.parseJson(json, ".parent.child.value"), (uint256));
+        assertEq(value, 42);
+
+        // Clean up the test file by removing the parent key we added
+        vm.removeFile(path);
+        vm.writeJson("{\"a\": 123, \"b\": \"0x000000000000000000000000000000000000bEEF\"}", path);
+    }
 }

testdata/cheats/Vm.sol

@@ -421,4 +421,28 @@ contract WriteTomlTest is DSTest {
         address decodedAddress = abi.decode(data, (address));
         assertEq(decodedAddress, ex);
     }
+
+    function test_writeToml_createKeys() public {
+        string memory path = "fixtures/Toml/write_test.toml";
+        string memory toml = vm.readFile(path);
+
+        bool exists = vm.keyExistsToml(toml, ".parent");
+        assertTrue(!exists);
+        exists = vm.keyExistsToml(toml, ".parent.child");
+        assertTrue(!exists);
+        exists = vm.keyExistsToml(toml, ".parent.child.value");
+        assertTrue(!exists);
+
+        // Write to nested path, creating intermediate keys
+        vm.writeToml(vm.toString(uint256(42)), path, ".parent.child.value");
+
+        // Verify the value was written and intermediate keys were created
+        toml = vm.readFile(path);
+        uint256 value = abi.decode(vm.parseToml(toml, ".parent.child.value"), (uint256));
+        assertEq(value, 42);
+
+        // Clean up the test file by removing the parent key we added
+        vm.removeFile(path);
+        vm.writeToml("{\"a\": 123, \"b\": \"0x000000000000000000000000000000000000bEEF\"}", path);
+    }
 }