Merge pull request #11 from initia-labs/feat/add-test-handler-for-cus…

…tom-query

feat: add test handler for custom query

Cargo.toml

@@ -100,6 +100,7 @@ bytes = "1.4.0"
 parking_lot = "0.12.1"
 base64 = "0.21.7"
 bigdecimal = "0.4"
+bech32 = "0.11"
 
 # Note: the BEGIN and END comments below are required for external tooling. Do not remove.
 # BEGIN MOVE DEPENDENCIES

crates/compiler/src/lib.rs

@@ -18,8 +18,8 @@ mod mocks;
 pub use clean::Clean;
 pub use command::Command;
 pub use compiler::execute;
-pub use new::New;
 pub use gas_meter::TestInitiaGasMeter;
+pub use new::New;
 
 #[cfg(test)]
 mod tests;

crates/e2e-move-tests/src/tests/move_unit.rs

@@ -11,7 +11,7 @@ use initia_move_natives::{
     account::NativeAccountContext, all_natives, block::NativeBlockContext, code::NativeCodeContext,
     cosmos::NativeCosmosContext, event::NativeEventContext, oracle::NativeOracleContext,
     staking::NativeStakingContext, table::NativeTableContext,
-    transaction_context::NativeTransactionContext,
+    transaction_context::NativeTransactionContext, query::NativeQueryContext,
 };
 use move_cli::base::test::{run_move_unit_tests_with_gas_meter, UnitTestResult};
 use move_core_types::effects::ChangeSet;
@@ -41,6 +41,7 @@ fn unit_test_extensions_hook(exts: &mut NativeContextExtensions) {
     exts.add(NativeTransactionContext::new([0; 32], [0; 32]));
     exts.add(NativeEventContext::default());
     exts.add(NativeOracleContext::new(&BLANK_API.oracle_api));
+    exts.add(NativeQueryContext::new(&BLANK_API.query_api));
 }
 
 fn run_tests_for_pkg(path_to_pkg: impl Into<String>) {

crates/natives/Cargo.toml

@@ -13,6 +13,7 @@ initia-move-gas = { workspace = true }
 
 anyhow = { workspace = true }
 bcs = { workspace = true }
+bech32 = { workspace = true }
 better_any = { workspace = true }
 serde = { workspace = true }
 serde_bytes = { workspace = true }

crates/natives/src/query.rs

@@ -1,16 +1,18 @@
+use bech32::{Bech32, Hrp};
 use better_any::{Tid, TidAble};
 use initia_move_gas::{InternalGas, GAS_UNIT_SCALING_FACTOR};
 use initia_move_types::query::*;
 use move_binary_format::errors::PartialVMError;
-use move_core_types::gas_algebra::NumBytes;
 use move_core_types::vm_status::StatusCode;
+use move_core_types::{account_address::AccountAddress, gas_algebra::NumBytes};
 use move_vm_runtime::native_functions::NativeFunction;
 use move_vm_types::{
     loaded_data::runtime_types::Type,
     values::{Value, Vector},
 };
+use serde::{Deserialize, Serialize};
 use smallvec::{smallvec, SmallVec};
-use std::collections::VecDeque;
+use std::collections::{BTreeMap, VecDeque};
 
 use crate::{
     interface::{
@@ -41,14 +43,22 @@ pub trait QueryAPI {
 #[derive(Tid)]
 pub struct NativeQueryContext<'a> {
     api: &'a dyn QueryAPI,
+
+    #[cfg(feature = "testing")]
+    responses: BTreeMap<String, Vec<u8>>,
 }
 
 /// Implementation of Native Query Context
 impl<'a> NativeQueryContext<'a> {
     /// Create a new instance of a native query context. This must be passed in via an
     /// extension into VM session functions.
     pub fn new(api: &'a dyn QueryAPI) -> Self {
-        Self { api }
+        Self {
+            api,
+
+            #[cfg(feature = "testing")]
+            responses: BTreeMap::new(),
+        }
     }
 }
 
@@ -75,13 +85,41 @@ fn native_query_custom(
         abort_code: UNABLE_TO_PARSE_STRING,
     })?;
 
+    let query_context = context.extensions().get::<NativeQueryContext>();
+
+    #[cfg(feature = "testing")]
+    if !name.is_empty() {
+        match name.as_str() {
+            "to_sdk_address" => {
+                return to_sdk_address(&data);
+            }
+            "from_sdk_address" => {
+                return from_sdk_address(&data);
+            }
+            _ => {
+                let mut hasher = Sha3_256::new();
+                hasher.update(name.as_bytes());
+                hasher.update(&data);
+                let hash = hex::encode(hasher.finalize());
+
+                let res = query_context
+                    .responses
+                    .get(&hash)
+                    .ok_or(SafeNativeError::Abort {
+                        abort_code: UNKNOWN_QUERY,
+                    })?;
+
+                return Ok(smallvec![Value::vector_u8(res.clone())]);
+            }
+        }
+    }
+
     let custom_query = CustomQuery { name, data };
     let req = QueryRequest::Custom(custom_query);
     let req = serde_json::to_vec(&req)
         .map_err(|err| partial_error(StatusCode::VALUE_SERIALIZATION_ERROR, err))?;
 
     let gas_balance: u64 = context.gas_balance().into();
-    let query_context = context.extensions().get::<NativeQueryContext>();
     let (res, used_gas) = query_context
         .api
         .query(req.as_slice(), gas_balance / GAS_UNIT_SCALING_FACTOR);
@@ -124,13 +162,31 @@ fn native_query_stargate(
         abort_code: UNABLE_TO_PARSE_STRING,
     })?;
 
+    let query_context = context.extensions().get::<NativeQueryContext>();
+
+    #[cfg(feature = "testing")]
+    if !path.is_empty() {
+        let mut hasher = Sha3_256::new();
+        hasher.update(path.as_bytes());
+        hasher.update(&data);
+        let hash = hex::encode(hasher.finalize());
+
+        let res = query_context
+            .responses
+            .get(&hash)
+            .ok_or(SafeNativeError::Abort {
+                abort_code: UNKNOWN_QUERY,
+            })?;
+
+        return Ok(smallvec![Value::vector_u8(res.clone())]);
+    }
+
     let stargate_query = StargateQuery { path, data };
     let req = QueryRequest::Stargate(stargate_query);
     let req = serde_json::to_vec(&req)
         .map_err(|err| partial_error(StatusCode::VALUE_SERIALIZATION_ERROR, err))?;
 
     let gas_balance: u64 = context.gas_balance().into();
-    let query_context = context.extensions().get::<NativeQueryContext>();
     let (res, used_gas) = query_context
         .api
         .query(req.as_slice(), gas_balance / GAS_UNIT_SCALING_FACTOR);
@@ -156,6 +212,8 @@ pub fn make_all(
     let natives = vec![
         ("query_custom", native_query_custom as RawSafeNative),
         ("query_stargate", native_query_stargate),
+        #[cfg(feature = "testing")]
+        ("set_query_response", native_test_only_set_query_response),
     ];
 
     builder.make_named_natives(natives)
@@ -167,3 +225,108 @@ pub fn make_all(
 fn partial_error(code: StatusCode, msg: impl ToString) -> PartialVMError {
     PartialVMError::new(code).with_message(msg.to_string())
 }
+
+#[cfg(feature = "testing")]
+use sha3::{Digest, Sha3_256};
+
+#[cfg(feature = "testing")]
+const UNKNOWN_QUERY: u64 = (ECATEGORY_INVALID_ARGUMENT << 16) + 110;
+
+#[cfg(feature = "testing")]
+#[derive(Deserialize)]
+struct ToSDKAddressRequest {
+    vm_addr: String,
+}
+
+#[cfg(feature = "testing")]
+#[derive(Serialize)]
+struct ToSDKAddressResponse {
+    sdk_addr: String,
+}
+
+#[cfg(feature = "testing")]
+#[derive(Deserialize)]
+struct FromSDKAddressRequest {
+    sdk_addr: String,
+}
+
+#[cfg(feature = "testing")]
+#[derive(Serialize)]
+struct FromSDKAddressResponse {
+    vm_addr: String,
+}
+
+#[cfg(feature = "testing")]
+fn to_sdk_address(data: &[u8]) -> SafeNativeResult<SmallVec<[Value; 1]>> {
+    let req: ToSDKAddressRequest =
+        serde_json::from_slice(data).map_err(|_| SafeNativeError::Abort {
+            abort_code: UNABLE_TO_PARSE_STRING,
+        })?;
+    let vm_addr =
+        AccountAddress::from_hex_literal(&req.vm_addr).map_err(|_| SafeNativeError::Abort {
+            abort_code: UNABLE_TO_PARSE_STRING,
+        })?;
+    let mut addr_bytes = vm_addr.as_slice();
+    if addr_bytes.starts_with(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) {
+        addr_bytes = &addr_bytes[12..];
+    }
+
+    let sdk_addr = bech32::encode::<Bech32>(Hrp::parse_unchecked("init"), addr_bytes).unwrap();
+    let res_bytes = serde_json::to_vec(&ToSDKAddressResponse { sdk_addr }).unwrap();
+    Ok(smallvec![Value::vector_u8(res_bytes)])
+}
+
+#[cfg(feature = "testing")]
+fn from_sdk_address(data: &[u8]) -> SafeNativeResult<SmallVec<[Value; 1]>> {
+    let req: FromSDKAddressRequest =
+        serde_json::from_slice(data).map_err(|_| SafeNativeError::Abort {
+            abort_code: UNABLE_TO_PARSE_STRING,
+        })?;
+    let (_, mut addr_bytes) =
+        bech32::decode(&req.sdk_addr).map_err(|_| SafeNativeError::Abort {
+            abort_code: UNABLE_TO_PARSE_STRING,
+        })?;
+
+    if addr_bytes.len() < AccountAddress::LENGTH {
+        let mut zero_padding = vec![0u8; AccountAddress::LENGTH - addr_bytes.len()];
+        zero_padding.append(&mut addr_bytes);
+
+        addr_bytes = zero_padding;
+    }
+
+    let vm_addr = AccountAddress::from_bytes(addr_bytes).unwrap();
+    let res_bytes = serde_json::to_vec(&FromSDKAddressResponse {
+        vm_addr: vm_addr.to_canonical_string(),
+    })
+    .unwrap();
+    Ok(smallvec![Value::vector_u8(res_bytes)])
+}
+
+#[cfg(feature = "testing")]
+fn native_test_only_set_query_response(
+    context: &mut SafeNativeContext,
+    ty_args: Vec<Type>,
+    mut arguments: VecDeque<Value>,
+) -> SafeNativeResult<SmallVec<[Value; 1]>> {
+    debug_assert!(ty_args.is_empty());
+    debug_assert!(arguments.len() == 3);
+
+    let output_bytes = safely_pop_arg!(arguments, Vector).to_vec_u8()?;
+    let data = safely_pop_arg!(arguments, Vector).to_vec_u8()?;
+    let path_or_name_bytes = safely_pop_arg!(arguments, Vector).to_vec_u8()?;
+
+    let path_or_name =
+        String::from_utf8(path_or_name_bytes).map_err(|_| SafeNativeError::Abort {
+            abort_code: UNABLE_TO_PARSE_STRING,
+        })?;
+
+    let mut hasher = Sha3_256::new();
+    hasher.update(path_or_name.as_bytes());
+    hasher.update(&data);
+    let hash = hex::encode(hasher.finalize());
+
+    let query_context = context.extensions_mut().get_mut::<NativeQueryContext>();
+    query_context.responses.insert(hash, output_bytes);
+
+    Ok(smallvec![])
+}

precompile/modules/initia_stdlib/sources/address.move

@@ -53,6 +53,20 @@ module initia_std::address {
         assert!(from_string(addr_str) == addr, 0)
     }
 
+    #[test]
+    fun test_to_sdk() {
+        let addr = @0x123abc;
+        let addr_sdk = string::utf8(b"init1qqqqqqqqqqqqqqqqqqqqqqqqqqqpyw4utfmfp0");
+        assert!(to_sdk(addr) == addr_sdk, 0)
+    }
+
+    #[test]
+    fun test_from_sdk() {
+        let addr = @0x123abc;
+        let addr_sdk = string::utf8(b"init1qqqqqqqqqqqqqqqqqqqqqqqqqqqpyw4utfmfp0");
+        assert!(addr == from_sdk(addr_sdk), 0)
+    }
+
     public native fun to_string(addr: address): String;
     public native fun from_string(addr_str: String): address;
 }

precompile/modules/initia_stdlib/sources/query.move

@@ -97,4 +97,23 @@ module initia_std::query {
     /// query_custom examples are in initia_stdlib::address module
     native public fun query_custom(name: vector<u8>, data: vector<u8>): vector<u8>;
     native public fun query_stargate(path: vector<u8>, data: vector<u8>): vector<u8>;
+
+    #[test_only]
+    native public fun set_query_response(path_or_name: vector<u8>, data: vector<u8>, response: vector<u8>);
+
+    #[test]
+    fun test_query_custom() {
+        set_query_response(b"path", b"data123", b"output");
+
+        let res = query_custom(b"path", b"data123");
+        assert!(res == b"output", 0);
+    }
+
+    #[test]
+    fun test_query_stargate() {
+        set_query_response(b"path", b"data123", b"output");
+
+        let res = query_stargate(b"path", b"data123");
+        assert!(res == b"output", 0);
+    }
 }

precompile/modules/minitia_stdlib/sources/address.move

@@ -53,6 +53,20 @@ module minitia_std::address {
         assert!(from_string(addr_str) == addr, 0)
     }
 
+    #[test]
+    fun test_to_sdk() {
+        let addr = @0x123abc;
+        let addr_sdk = string::utf8(b"init1qqqqqqqqqqqqqqqqqqqqqqqqqqqpyw4utfmfp0");
+        assert!(to_sdk(addr) == addr_sdk, 0)
+    }
+
+    #[test]
+    fun test_from_sdk() {
+        let addr = @0x123abc;
+        let addr_sdk = string::utf8(b"init1qqqqqqqqqqqqqqqqqqqqqqqqqqqpyw4utfmfp0");
+        assert!(addr == from_sdk(addr_sdk), 0)
+    }
+
     public native fun to_string(addr: address): String;
     public native fun from_string(addr_str: String): address;
 }

precompile/modules/minitia_stdlib/sources/query.move

@@ -97,4 +97,23 @@ module minitia_std::query {
     /// query_custom examples are in initia_stdlib::address module
     native public fun query_custom(name: vector<u8>, data: vector<u8>): vector<u8>;
     native public fun query_stargate(path: vector<u8>, data: vector<u8>): vector<u8>;
+
+    #[test_only]
+    native public fun set_query_response(path_or_name: vector<u8>, data: vector<u8>, response: vector<u8>);
+
+    #[test]
+    fun test_query_custom() {
+        set_query_response(b"path", b"data123", b"output");
+
+        let res = query_custom(b"path", b"data123");
+        assert!(res == b"output", 0);
+    }
+
+    #[test]
+    fun test_query_stargate() {
+        set_query_response(b"path", b"data123", b"output");
+
+        let res = query_stargate(b"path", b"data123");
+        assert!(res == b"output", 0);
+    }
 }