Add parseJsonKeys() cheatcode (#5252)

* add fetchJsonKeys() cheatcode

* fmt

* forge fmt

* fix tests

* Remove unwrap()

* Fix typo

abi/abi/HEVM.sol

@@ -180,6 +180,7 @@ writeJson(string, string)
 writeJson(string, string, string)
 parseJson(string)(bytes)
 parseJson(string, string)(bytes)
+parseJsonKeys(string, string)(string[])
 parseJsonUint(string, string)(uint256)
 parseJsonUintArray(string, string)(uint256[])
 parseJsonInt(string, string)(int256)

evm/src/executor/inspector/cheatcodes/ext.rs

@@ -319,6 +319,36 @@ fn parse_json(json_str: &str, key: &str, coerce: Option<ParamType>) -> Result {
     }
 }
 
+// returns JSON keys of given object as a string array
+fn parse_json_keys(json_str: &str, key: &str) -> Result {
+    let json = serde_json::from_str(json_str)?;
+    let values = jsonpath_lib::select(&json, &canonicalize_json_key(key))?;
+
+    // We need to check that values contains just one JSON-object and not an array of objects
+    ensure!(
+        values.len() == 1,
+        "You can only get keys for a single JSON-object. The key '{key}' returns a value or an array of JSON-objects",
+    );
+
+    let value = values[0];
+
+    ensure!(
+        value.is_object(),
+        "You can only get keys for JSON-object. The key '{key}' does not return an object",
+    );
+
+    let res = value
+        .as_object()
+        .ok_or(eyre::eyre!("Unexpected error while extracting JSON-object"))?
+        .keys()
+        .map(|key| Token::String(key.to_owned()))
+        .collect::<Vec<Token>>();
+
+    // encode the bytes as the 'bytes' solidity type
+    let abi_encoded = abi::encode(&[Token::Array(res)]);
+    Ok(abi_encoded.into())
+}
+
 /// Serializes a key:value pair to a specific object. By calling this function multiple times,
 /// the user can serialize multiple KV pairs to the same object. The value can be of any type, even
 /// a new object in itself. The function will return
@@ -523,6 +553,7 @@ pub fn apply(state: &mut Cheatcodes, call: &HEVMCalls) -> Option<Result> {
         HEVMCalls::ParseJson0(inner) => parse_json(&inner.0, "$", None),
         HEVMCalls::ParseJson1(inner) => parse_json(&inner.0, &inner.1, None),
         HEVMCalls::ParseJsonBool(inner) => parse_json(&inner.0, &inner.1, Some(ParamType::Bool)),
+        HEVMCalls::ParseJsonKeys(inner) => parse_json_keys(&inner.0, &inner.1),
         HEVMCalls::ParseJsonBoolArray(inner) => {
             parse_json(&inner.0, &inner.1, Some(ParamType::Bool))
         }

testdata/cheats/Json.t.sol

@@ -159,6 +159,27 @@ contract ParseJsonTest is DSTest {
         data = vm.parseJson("", ".");
         assertEq(0, data.length);
     }
+
+    function test_parseJsonKeys() public {
+        string memory jsonString =
+            '{"some_key_to_value": "some_value", "some_key_to_array": [1, 2, 3], "some_key_to_object": {"key1": "value1", "key2": 2}}';
+        string[] memory keys = vm.parseJsonKeys(jsonString, "$");
+        assertEq(abi.encode(keys), abi.encode(["some_key_to_value", "some_key_to_array", "some_key_to_object"]));
+
+        keys = vm.parseJsonKeys(jsonString, ".some_key_to_object");
+        assertEq(abi.encode(keys), abi.encode(["key1", "key2"]));
+
+        vm.expectRevert("You can only get keys for JSON-object. The key '.some_key_to_array' does not return an object");
+        vm.parseJsonKeys(jsonString, ".some_key_to_array");
+
+        vm.expectRevert("You can only get keys for JSON-object. The key '.some_key_to_value' does not return an object");
+        vm.parseJsonKeys(jsonString, ".some_key_to_value");
+
+        vm.expectRevert(
+            "You can only get keys for a single JSON-object. The key '.*' returns a value or an array of JSON-objects"
+        );
+        vm.parseJsonKeys(jsonString, ".*");
+    }
 }
 
 contract WriteJsonTest is DSTest {

testdata/cheats/Vm.sol

@@ -525,6 +525,8 @@ interface Vm {
 
     function parseJson(string calldata) external returns (bytes memory);
 
+    function parseJsonKeys(string calldata, string calldata) external returns (string[] memory);
+
     function parseJsonUint(string calldata, string calldata) external returns (uint256);
 
     function parseJsonUintArray(string calldata, string calldata) external returns (uint256[] memory);