Handle reply with reflect contract

CosmWasm · Jul 13, 2021 · d81b758 · d81b758
1 parent 7e71ddd
commit d81b758
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 13 deletions.
diff --git a/packages/multi-test/src/app.rs b/packages/multi-test/src/app.rs
@@ -405,6 +405,7 @@ mod test {
     use crate::test_helpers::{
         contract_payout, contract_payout_custom, contract_reflect, CustomMsg, EmptyMsg,
         PayoutCountResponse, PayoutInitMessage, PayoutQueryMsg, PayoutSudoMsg, ReflectMessage,
+        ReflectQueryMsg,
     };
     use crate::SimpleBank;
     use cosmwasm_std::testing::MockStorage;
@@ -547,7 +548,7 @@ mod test {
         // reflect count is 1
         let qres: PayoutCountResponse = router
             .wrap()
-            .query_wasm_smart(&reflect_addr, &EmptyMsg {})
+            .query_wasm_smart(&reflect_addr, &ReflectQueryMsg::Count {})
             .unwrap();
         assert_eq!(0, qres.count);
 
@@ -575,7 +576,7 @@ mod test {
         // reflect count updated
         let qres: PayoutCountResponse = router
             .wrap()
-            .query_wasm_smart(&reflect_addr, &EmptyMsg {})
+            .query_wasm_smart(&reflect_addr, &ReflectQueryMsg::Count {})
             .unwrap();
         assert_eq!(1, qres.count);
     }
@@ -622,10 +623,10 @@ mod test {
         let funds = get_balance(&router, &random);
         assert_eq!(funds, coins(7, "eth"));
 
-        // reflect count should be updated to 2
+        // reflect count should be updated to 1
         let qres: PayoutCountResponse = router
             .wrap()
-            .query_wasm_smart(&reflect_addr, &EmptyMsg {})
+            .query_wasm_smart(&reflect_addr, &ReflectQueryMsg::Count {})
             .unwrap();
         assert_eq!(1, qres.count);
 
@@ -653,7 +654,7 @@ mod test {
         // failure should not update reflect count
         let qres: PayoutCountResponse = router
             .wrap()
-            .query_wasm_smart(&reflect_addr, &EmptyMsg {})
+            .query_wasm_smart(&reflect_addr, &ReflectQueryMsg::Count {})
             .unwrap();
         assert_eq!(1, qres.count);
     }

diff --git a/packages/multi-test/src/test_helpers.rs b/packages/multi-test/src/test_helpers.rs
@@ -4,10 +4,10 @@ use serde::{Deserialize, Serialize};
 use std::fmt;
 
 use cosmwasm_std::{
-    attr, to_binary, BankMsg, Binary, Coin, Deps, DepsMut, Empty, Env, MessageInfo, Response,
-    StdError, SubMsg,
+    attr, to_binary, BankMsg, Binary, Coin, Deps, DepsMut, Empty, Env, MessageInfo, Reply,
+    Response, StdError, SubMsg,
 };
-use cw_storage_plus::Item;
+use cw_storage_plus::{Item, Map, U64Key};
 
 use crate::wasm::{Contract, ContractWrapper};
 
@@ -159,6 +159,14 @@ pub struct ReflectMessage {
     pub messages: Vec<SubMsg<CustomMsg>>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum ReflectQueryMsg {
+    Count {},
+    Reply { id: u64 },
+}
+
+const REFLECT: Map<U64Key, Reply> = Map::new("reflect");
+
 fn instantiate_reflect(
     deps: DepsMut,
     _env: Env,
@@ -186,13 +194,31 @@ fn execute_reflect(
     Ok(res)
 }
 
-fn query_reflect(deps: Deps, _env: Env, _msg: EmptyMsg) -> Result<Binary, StdError> {
-    let count = COUNT.load(deps.storage)?;
-    let res = PayoutCountResponse { count };
-    to_binary(&res)
+fn query_reflect(deps: Deps, _env: Env, msg: ReflectQueryMsg) -> Result<Binary, StdError> {
+    match msg {
+        ReflectQueryMsg::Count {} => {
+            let count = COUNT.load(deps.storage)?;
+            let res = PayoutCountResponse { count };
+            to_binary(&res)
+        }
+        ReflectQueryMsg::Reply { id } => {
+            let reply = REFLECT.load(deps.storage, id.into())?;
+            to_binary(&reply)
+        }
+    }
+}
+
+fn reply_reflect(deps: DepsMut, _env: Env, msg: Reply) -> Result<Response<CustomMsg>, StdError> {
+    REFLECT.save(deps.storage, msg.id.into(), &msg)?;
+    Ok(Response::default())
 }
 
 pub fn contract_reflect() -> Box<dyn Contract<CustomMsg>> {
-    let contract = ContractWrapper::new(execute_reflect, instantiate_reflect, query_reflect);
+    let contract = ContractWrapper::new_with_reply(
+        execute_reflect,
+        instantiate_reflect,
+        query_reflect,
+        reply_reflect,
+    );
     Box::new(contract)
 }