From 3edbdc75dfa3c3f72c2273a7488881d1362ac626 Mon Sep 17 00:00:00 2001
From: Tomasz Kurcz <uint@lavabit.com>
Date: Tue, 8 Feb 2022 22:25:35 +0100
Subject: [PATCH] multi-test: support custom queries

---
 packages/multi-test/src/app.rs       |   2 +-
 packages/multi-test/src/contracts.rs | 130 ++++++++++++++-------------
 packages/multi-test/src/wasm.rs      |   8 +-
 3 files changed, 75 insertions(+), 65 deletions(-)

diff --git a/packages/multi-test/src/app.rs b/packages/multi-test/src/app.rs
index c23e6873a..f06ba6ac4 100644
--- a/packages/multi-test/src/app.rs
+++ b/packages/multi-test/src/app.rs
@@ -557,7 +557,7 @@ where
 {
     /// This registers contract code (like uploading wasm bytecode on a chain),
     /// so it can later be used to instantiate a contract.
-    pub fn store_code(&mut self, code: Box<dyn Contract<CustomT::ExecT>>) -> u64 {
+    pub fn store_code(&mut self, code: Box<dyn Contract<CustomT::ExecT, CustomT::QueryT>>) -> u64 {
         self.init_modules(|router, _, _| router.wasm.store_code(code) as u64)
     }
 
diff --git a/packages/multi-test/src/contracts.rs b/packages/multi-test/src/contracts.rs
index e703b75a8..72bf76950 100644
--- a/packages/multi-test/src/contracts.rs
+++ b/packages/multi-test/src/contracts.rs
@@ -4,19 +4,21 @@ use std::error::Error;
 use std::fmt::{self, Debug, Display};
 
 use cosmwasm_std::{
-    from_slice, Binary, CosmosMsg, Deps, DepsMut, Empty, Env, MessageInfo, Reply, Response, SubMsg,
+    from_slice, Binary, CosmosMsg, CustomQuery, Deps, DepsMut, Empty, Env, MessageInfo, Reply,
+    Response, SubMsg,
 };
 
 use anyhow::{anyhow, bail, Context, Result as AnyResult};
 
 /// Interface to call into a Contract
-pub trait Contract<T>
+pub trait Contract<T, Q = Empty>
 where
     T: Clone + fmt::Debug + PartialEq + JsonSchema,
+    Q: CustomQuery,
 {
     fn execute(
         &self,
-        deps: DepsMut,
+        deps: DepsMut<Q>,
         env: Env,
         info: MessageInfo,
         msg: Vec<u8>,
@@ -24,31 +26,32 @@ where
 
     fn instantiate(
         &self,
-        deps: DepsMut,
+        deps: DepsMut<Q>,
         env: Env,
         info: MessageInfo,
         msg: Vec<u8>,
     ) -> AnyResult<Response<T>>;
 
-    fn query(&self, deps: Deps, env: Env, msg: Vec<u8>) -> AnyResult<Binary>;
+    fn query(&self, deps: Deps<Q>, env: Env, msg: Vec<u8>) -> AnyResult<Binary>;
 
-    fn sudo(&self, deps: DepsMut, env: Env, msg: Vec<u8>) -> AnyResult<Response<T>>;
+    fn sudo(&self, deps: DepsMut<Q>, env: Env, msg: Vec<u8>) -> AnyResult<Response<T>>;
 
-    fn reply(&self, deps: DepsMut, env: Env, msg: Reply) -> AnyResult<Response<T>>;
+    fn reply(&self, deps: DepsMut<Q>, env: Env, msg: Reply) -> AnyResult<Response<T>>;
 
-    fn migrate(&self, deps: DepsMut, env: Env, msg: Vec<u8>) -> AnyResult<Response<T>>;
+    fn migrate(&self, deps: DepsMut<Q>, env: Env, msg: Vec<u8>) -> AnyResult<Response<T>>;
 }
 
-type ContractFn<T, C, E> =
-    fn(deps: DepsMut, env: Env, info: MessageInfo, msg: T) -> Result<Response<C>, E>;
-type PermissionedFn<T, C, E> = fn(deps: DepsMut, env: Env, msg: T) -> Result<Response<C>, E>;
-type ReplyFn<C, E> = fn(deps: DepsMut, env: Env, msg: Reply) -> Result<Response<C>, E>;
-type QueryFn<T, E> = fn(deps: Deps, env: Env, msg: T) -> Result<Binary, E>;
+type ContractFn<T, C, E, Q> =
+    fn(deps: DepsMut<Q>, env: Env, info: MessageInfo, msg: T) -> Result<Response<C>, E>;
+type PermissionedFn<T, C, E, Q> = fn(deps: DepsMut<Q>, env: Env, msg: T) -> Result<Response<C>, E>;
+type ReplyFn<C, E, Q> = fn(deps: DepsMut<Q>, env: Env, msg: Reply) -> Result<Response<C>, E>;
+type QueryFn<T, E, Q> = fn(deps: Deps<Q>, env: Env, msg: T) -> Result<Binary, E>;
 
-type ContractClosure<T, C, E> = Box<dyn Fn(DepsMut, Env, MessageInfo, T) -> Result<Response<C>, E>>;
-type PermissionedClosure<T, C, E> = Box<dyn Fn(DepsMut, Env, T) -> Result<Response<C>, E>>;
-type ReplyClosure<C, E> = Box<dyn Fn(DepsMut, Env, Reply) -> Result<Response<C>, E>>;
-type QueryClosure<T, E> = Box<dyn Fn(Deps, Env, T) -> Result<Binary, E>>;
+type ContractClosure<T, C, E, Q> =
+    Box<dyn Fn(DepsMut<Q>, Env, MessageInfo, T) -> Result<Response<C>, E>>;
+type PermissionedClosure<T, C, E, Q> = Box<dyn Fn(DepsMut<Q>, Env, T) -> Result<Response<C>, E>>;
+type ReplyClosure<C, E, Q> = Box<dyn Fn(DepsMut<Q>, Env, Reply) -> Result<Response<C>, E>>;
+type QueryClosure<T, E, Q> = Box<dyn Fn(Deps<Q>, Env, T) -> Result<Binary, E>>;
 
 /// Wraps the exported functions from a contract and provides the normalized format
 /// Place T4 and E4 at the end, as we just want default placeholders for most contracts that don't have sudo
@@ -60,6 +63,7 @@ pub struct ContractWrapper<
     E2,
     E3,
     C = Empty,
+    Q = Empty,
     T4 = Empty,
     E4 = anyhow::Error,
     E5 = anyhow::Error,
@@ -78,16 +82,17 @@ pub struct ContractWrapper<
     E5: Display + Debug + Send + Sync + 'static,
     E6: Display + Debug + Send + Sync + 'static,
     C: Clone + fmt::Debug + PartialEq + JsonSchema,
+    Q: CustomQuery + DeserializeOwned + 'static,
 {
-    execute_fn: ContractClosure<T1, C, E1>,
-    instantiate_fn: ContractClosure<T2, C, E2>,
-    query_fn: QueryClosure<T3, E3>,
-    sudo_fn: Option<PermissionedClosure<T4, C, E4>>,
-    reply_fn: Option<ReplyClosure<C, E5>>,
-    migrate_fn: Option<PermissionedClosure<T6, C, E6>>,
+    execute_fn: ContractClosure<T1, C, E1, Q>,
+    instantiate_fn: ContractClosure<T2, C, E2, Q>,
+    query_fn: QueryClosure<T3, E3, Q>,
+    sudo_fn: Option<PermissionedClosure<T4, C, E4, Q>>,
+    reply_fn: Option<ReplyClosure<C, E5, Q>>,
+    migrate_fn: Option<PermissionedClosure<T6, C, E6, Q>>,
 }
 
-impl<T1, T2, T3, E1, E2, E3, C> ContractWrapper<T1, T2, T3, E1, E2, E3, C>
+impl<T1, T2, T3, E1, E2, E3, C, Q> ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q>
 where
     T1: DeserializeOwned + Debug + 'static,
     T2: DeserializeOwned + 'static,
@@ -96,13 +101,14 @@ where
     E2: Display + Debug + Send + Sync + 'static,
     E3: Display + Debug + Send + Sync + 'static,
     C: Clone + fmt::Debug + PartialEq + JsonSchema + 'static,
+    Q: CustomQuery + DeserializeOwned + 'static,
 {
     pub fn new(
-        execute_fn: ContractFn<T1, C, E1>,
-        instantiate_fn: ContractFn<T2, C, E2>,
-        query_fn: QueryFn<T3, E3>,
+        execute_fn: ContractFn<T1, C, E1, Q>,
+        instantiate_fn: ContractFn<T2, C, E2, Q>,
+        query_fn: QueryFn<T3, E3, Q>,
     ) -> Self {
-        ContractWrapper {
+        Self {
             execute_fn: Box::new(execute_fn),
             instantiate_fn: Box::new(instantiate_fn),
             query_fn: Box::new(query_fn),
@@ -115,11 +121,11 @@ where
     /// this will take a contract that returns Response<Empty> and will "upgrade" it
     /// to Response<C> if needed to be compatible with a chain-specific extension
     pub fn new_with_empty(
-        execute_fn: ContractFn<T1, Empty, E1>,
-        instantiate_fn: ContractFn<T2, Empty, E2>,
-        query_fn: QueryFn<T3, E3>,
+        execute_fn: ContractFn<T1, Empty, E1, Q>,
+        instantiate_fn: ContractFn<T2, Empty, E2, Q>,
+        query_fn: QueryFn<T3, E3, Q>,
     ) -> Self {
-        ContractWrapper {
+        Self {
             execute_fn: customize_fn(execute_fn),
             instantiate_fn: customize_fn(instantiate_fn),
             query_fn: Box::new(query_fn),
@@ -130,8 +136,8 @@ where
     }
 }
 
-impl<T1, T2, T3, E1, E2, E3, C, T4, E4, E5, T6, E6>
-    ContractWrapper<T1, T2, T3, E1, E2, E3, C, T4, E4, E5, T6, E6>
+impl<T1, T2, T3, E1, E2, E3, C, Q, T4, E4, E5, T6, E6>
+    ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q, T4, E4, E5, T6, E6>
 where
     T1: DeserializeOwned + Debug + 'static,
     T2: DeserializeOwned + 'static,
@@ -145,11 +151,12 @@ where
     E5: Display + Debug + Send + Sync + 'static,
     E6: Display + Debug + Send + Sync + 'static,
     C: Clone + fmt::Debug + PartialEq + JsonSchema + 'static,
+    Q: CustomQuery + DeserializeOwned + 'static,
 {
     pub fn with_sudo<T4A, E4A>(
         self,
-        sudo_fn: PermissionedFn<T4A, C, E4A>,
-    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, T4A, E4A, E5, T6, E6>
+        sudo_fn: PermissionedFn<T4A, C, E4A, Q>,
+    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q, T4A, E4A, E5, T6, E6>
     where
         T4A: DeserializeOwned + 'static,
         E4A: Display + Debug + Send + Sync + 'static,
@@ -166,8 +173,8 @@ where
 
     pub fn with_sudo_empty<T4A, E4A>(
         self,
-        sudo_fn: PermissionedFn<T4A, Empty, E4A>,
-    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, T4A, E4A, E5, T6, E6>
+        sudo_fn: PermissionedFn<T4A, Empty, E4A, Q>,
+    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q, T4A, E4A, E5, T6, E6>
     where
         T4A: DeserializeOwned + 'static,
         E4A: Display + Debug + Send + Sync + 'static,
@@ -184,8 +191,8 @@ where
 
     pub fn with_reply<E5A>(
         self,
-        reply_fn: ReplyFn<C, E5A>,
-    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, T4, E4, E5A, T6, E6>
+        reply_fn: ReplyFn<C, E5A, Q>,
+    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q, T4, E4, E5A, T6, E6>
     where
         E5A: Display + Debug + Send + Sync + 'static,
     {
@@ -202,8 +209,8 @@ where
     /// A correlate of new_with_empty
     pub fn with_reply_empty<E5A>(
         self,
-        reply_fn: ReplyFn<Empty, E5A>,
-    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, T4, E4, E5A, T6, E6>
+        reply_fn: ReplyFn<Empty, E5A, Q>,
+    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q, T4, E4, E5A, T6, E6>
     where
         E5A: Display + Debug + Send + Sync + 'static,
     {
@@ -219,8 +226,8 @@ where
 
     pub fn with_migrate<T6A, E6A>(
         self,
-        migrate_fn: PermissionedFn<T6A, C, E6A>,
-    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, T4, E4, E5, T6A, E6A>
+        migrate_fn: PermissionedFn<T6A, C, E6A, Q>,
+    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q, T4, E4, E5, T6A, E6A>
     where
         T6A: DeserializeOwned + 'static,
         E6A: Display + Debug + Send + Sync + 'static,
@@ -237,8 +244,8 @@ where
 
     pub fn with_migrate_empty<T6A, E6A>(
         self,
-        migrate_fn: PermissionedFn<T6A, Empty, E6A>,
-    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, T4, E4, E5, T6A, E6A>
+        migrate_fn: PermissionedFn<T6A, Empty, E6A, Q>,
+    ) -> ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q, T4, E4, E5, T6A, E6A>
     where
         T6A: DeserializeOwned + 'static,
         E6A: Display + Debug + Send + Sync + 'static,
@@ -254,28 +261,30 @@ where
     }
 }
 
-fn customize_fn<T, C, E>(raw_fn: ContractFn<T, Empty, E>) -> ContractClosure<T, C, E>
+fn customize_fn<T, C, E, Q>(raw_fn: ContractFn<T, Empty, E, Q>) -> ContractClosure<T, C, E, Q>
 where
     T: DeserializeOwned + 'static,
     E: Display + Debug + Send + Sync + 'static,
     C: Clone + fmt::Debug + PartialEq + JsonSchema + 'static,
+    Q: CustomQuery + DeserializeOwned + 'static,
 {
     let customized =
-        move |deps: DepsMut, env: Env, info: MessageInfo, msg: T| -> Result<Response<C>, E> {
+        move |deps: DepsMut<Q>, env: Env, info: MessageInfo, msg: T| -> Result<Response<C>, E> {
             raw_fn(deps, env, info, msg).map(customize_response::<C>)
         };
     Box::new(customized)
 }
 
-fn customize_permissioned_fn<T, C, E>(
-    raw_fn: PermissionedFn<T, Empty, E>,
-) -> PermissionedClosure<T, C, E>
+fn customize_permissioned_fn<T, C, E, Q>(
+    raw_fn: PermissionedFn<T, Empty, E, Q>,
+) -> PermissionedClosure<T, C, E, Q>
 where
     T: DeserializeOwned + 'static,
     E: Display + Debug + Send + Sync + 'static,
     C: Clone + fmt::Debug + PartialEq + JsonSchema + 'static,
+    Q: CustomQuery + DeserializeOwned + 'static,
 {
-    let customized = move |deps: DepsMut, env: Env, msg: T| -> Result<Response<C>, E> {
+    let customized = move |deps: DepsMut<Q>, env: Env, msg: T| -> Result<Response<C>, E> {
         raw_fn(deps, env, msg).map(customize_response::<C>)
     };
     Box::new(customized)
@@ -315,8 +324,8 @@ where
     }
 }
 
-impl<T1, T2, T3, E1, E2, E3, C, T4, E4, E5, T6, E6> Contract<C>
-    for ContractWrapper<T1, T2, T3, E1, E2, E3, C, T4, E4, E5, T6, E6>
+impl<T1, T2, T3, E1, E2, E3, C, T4, E4, E5, T6, E6, Q> Contract<C, Q>
+    for ContractWrapper<T1, T2, T3, E1, E2, E3, C, Q, T4, E4, E5, T6, E6>
 where
     T1: DeserializeOwned + Debug + Clone,
     T2: DeserializeOwned + Debug + Clone,
@@ -330,10 +339,11 @@ where
     E5: Display + Debug + Send + Sync + 'static,
     E6: Display + Debug + Send + Sync + 'static,
     C: Clone + fmt::Debug + PartialEq + JsonSchema,
+    Q: CustomQuery + DeserializeOwned,
 {
     fn execute(
         &self,
-        deps: DepsMut,
+        deps: DepsMut<Q>,
         env: Env,
         info: MessageInfo,
         msg: Vec<u8>,
@@ -349,7 +359,7 @@ where
 
     fn instantiate(
         &self,
-        deps: DepsMut,
+        deps: DepsMut<Q>,
         env: Env,
         info: MessageInfo,
         msg: Vec<u8>,
@@ -363,7 +373,7 @@ where
             ))
     }
 
-    fn query(&self, deps: Deps, env: Env, msg: Vec<u8>) -> AnyResult<Binary> {
+    fn query(&self, deps: Deps<Q>, env: Env, msg: Vec<u8>) -> AnyResult<Binary> {
         let msg: T3 = from_slice(&msg)?;
         (self.query_fn)(deps, env, msg.clone())
             .map_err(anyhow::Error::from)
@@ -374,7 +384,7 @@ where
     }
 
     // this returns an error if the contract doesn't implement sudo
-    fn sudo(&self, deps: DepsMut, env: Env, msg: Vec<u8>) -> AnyResult<Response<C>> {
+    fn sudo(&self, deps: DepsMut<Q>, env: Env, msg: Vec<u8>) -> AnyResult<Response<C>> {
         let msg = from_slice(&msg)?;
         match &self.sudo_fn {
             Some(sudo) => sudo(deps, env, msg).map_err(|err| anyhow!(err)),
@@ -383,7 +393,7 @@ where
     }
 
     // this returns an error if the contract doesn't implement reply
-    fn reply(&self, deps: DepsMut, env: Env, reply_data: Reply) -> AnyResult<Response<C>> {
+    fn reply(&self, deps: DepsMut<Q>, env: Env, reply_data: Reply) -> AnyResult<Response<C>> {
         match &self.reply_fn {
             Some(reply) => reply(deps, env, reply_data).map_err(|err| anyhow!(err)),
             None => bail!("reply not implemented for contract"),
@@ -391,7 +401,7 @@ where
     }
 
     // this returns an error if the contract doesn't implement migrate
-    fn migrate(&self, deps: DepsMut, env: Env, msg: Vec<u8>) -> AnyResult<Response<C>> {
+    fn migrate(&self, deps: DepsMut<Q>, env: Env, msg: Vec<u8>) -> AnyResult<Response<C>> {
         let msg = from_slice(&msg)?;
         match &self.migrate_fn {
             Some(migrate) => migrate(deps, env, msg).map_err(|err| anyhow!(err)),
diff --git a/packages/multi-test/src/wasm.rs b/packages/multi-test/src/wasm.rs
index 9a30d47e7..ce793c3af 100644
--- a/packages/multi-test/src/wasm.rs
+++ b/packages/multi-test/src/wasm.rs
@@ -99,7 +99,7 @@ pub trait Wasm<ExecC, QueryC> {
 pub struct WasmKeeper<ExecC, QueryC> {
     /// code is in-memory lookup that stands in for wasm code
     /// this can only be edited on the WasmRouter, and just read in caches
-    codes: HashMap<usize, Box<dyn Contract<ExecC>>>,
+    codes: HashMap<usize, Box<dyn Contract<ExecC, QueryC>>>,
     /// Just markers to make type elision fork when using it as `Wasm` trait
     _p: std::marker::PhantomData<QueryC>,
 }
@@ -180,7 +180,7 @@ where
 }
 
 impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
-    pub fn store_code(&mut self, code: Box<dyn Contract<ExecC>>) -> usize {
+    pub fn store_code(&mut self, code: Box<dyn Contract<ExecC, QueryC>>) -> usize {
         let idx = self.codes.len() + 1;
         self.codes.insert(idx, code);
         idx
@@ -706,7 +706,7 @@ where
         action: F,
     ) -> AnyResult<T>
     where
-        F: FnOnce(&Box<dyn Contract<ExecC>>, Deps, Env) -> AnyResult<T>,
+        F: FnOnce(&Box<dyn Contract<ExecC, QueryC>>, Deps<QueryC>, Env) -> AnyResult<T>,
     {
         let contract = self.load_contract(storage, &address)?;
         let handler = self
@@ -734,7 +734,7 @@ where
         action: F,
     ) -> AnyResult<T>
     where
-        F: FnOnce(&Box<dyn Contract<ExecC>>, DepsMut, Env) -> AnyResult<T>,
+        F: FnOnce(&Box<dyn Contract<ExecC, QueryC>>, DepsMut<QueryC>, Env) -> AnyResult<T>,
         ExecC: DeserializeOwned,
     {
         let contract = self.load_contract(storage, &address)?;