Fix instaniate reply data

packages/multi-test/src/app.rs

@@ -2390,4 +2390,148 @@ mod test {
             assert!(custom_handler_state.queries().is_empty());
         }
     }
+
+    mod protobuf_wrapped_data {
+        use super::*;
+        use cw0::{parse_execute_response_data, parse_instantiate_response_data};
+
+        #[test]
+        fn instantiate_wrapped_properly() {
+            // set personal balance
+            let owner = Addr::unchecked("owner");
+            let init_funds = vec![coin(20, "btc")];
+
+            let mut app = custom_app::<CustomMsg, Empty, _>(|router, _, storage| {
+                router
+                    .bank
+                    .init_balance(storage, &owner, init_funds)
+                    .unwrap();
+            });
+
+            // set up reflect contract
+            let code_id = app.store_code(reflect::contract());
+            let init_msg = to_binary(&EmptyMsg {}).unwrap();
+            let msg = WasmMsg::Instantiate {
+                admin: None,
+                code_id,
+                msg: init_msg,
+                funds: vec![],
+                label: "label".into(),
+            };
+            let res = app.execute(owner, msg.into()).unwrap();
+
+            // assert we have a proper instantiate result
+            let parsed = parse_instantiate_response_data(res.data.unwrap().as_slice()).unwrap();
+            assert!(parsed.data.is_none());
+            // check the address is right
+
+            let count: payout::CountResponse = app
+                .wrap()
+                .query_wasm_smart(&parsed.contract_address, &reflect::QueryMsg::Count {})
+                .unwrap();
+            assert_eq!(count.count, 0);
+        }
+
+        #[test]
+        fn instantiate_with_data_works() {
+            let owner = Addr::unchecked("owner");
+            let mut app = BasicApp::new(|_, _, _| {});
+
+            // set up echo contract
+            let code_id = app.store_code(echo::contract());
+            let msg = echo::InitMessage::<Empty> {
+                data: Some("food".into()),
+                sub_msg: None,
+            };
+            let init_msg = to_binary(&msg).unwrap();
+            let msg = WasmMsg::Instantiate {
+                admin: None,
+                code_id,
+                msg: init_msg,
+                funds: vec![],
+                label: "label".into(),
+            };
+            let res = app.execute(owner, msg.into()).unwrap();
+
+            // assert we have a proper instantiate result
+            let parsed = parse_instantiate_response_data(res.data.unwrap().as_slice()).unwrap();
+            assert!(parsed.data.is_some());
+            assert_eq!(parsed.data.unwrap(), Binary::from(b"food"));
+            assert!(!parsed.contract_address.is_empty());
+        }
+
+        #[test]
+        fn instantiate_with_reply_works() {
+            let owner = Addr::unchecked("owner");
+            let mut app = BasicApp::new(|_, _, _| {});
+
+            // set up echo contract
+            let code_id = app.store_code(echo::contract());
+            let msg = echo::InitMessage::<Empty> {
+                data: Some("food".into()),
+                ..Default::default()
+            };
+            let addr1 = app
+                .instantiate_contract(code_id, owner.clone(), &msg, &[], "first", None)
+                .unwrap();
+
+            // another echo contract
+            let msg = echo::Message::<Empty> {
+                data: Some("babble".into()),
+                ..Default::default()
+            };
+            let sub_msg = SubMsg::reply_on_success(
+                WasmMsg::Execute {
+                    contract_addr: addr1.to_string(),
+                    msg: to_binary(&msg).unwrap(),
+                    funds: vec![],
+                },
+                1234,
+            );
+            let init_msg = echo::InitMessage::<Empty> {
+                data: Some("remove_me".into()),
+                sub_msg: Some(vec![sub_msg]),
+            };
+            let init_msg = to_binary(&init_msg).unwrap();
+            let msg = WasmMsg::Instantiate {
+                admin: None,
+                code_id,
+                msg: init_msg,
+                funds: vec![],
+                label: "label".into(),
+            };
+            let res = app.execute(owner, msg.into()).unwrap();
+
+            // assert we have a proper instantiate result
+            let parsed = parse_instantiate_response_data(res.data.unwrap().as_slice()).unwrap();
+            assert!(parsed.data.is_some());
+            // from the reply, not the top level
+            assert_eq!(parsed.data.unwrap(), Binary::from(b"babble"));
+            assert!(!parsed.contract_address.is_empty());
+            assert_ne!(parsed.contract_address, addr1.to_string());
+        }
+
+        #[ignore]
+        #[test]
+        fn execute_wrapped_properly() {
+            let owner = Addr::unchecked("owner");
+            let mut app = BasicApp::new(|_, _, _| {});
+
+            // set up reflect contract
+            let code_id = app.store_code(echo::contract());
+            let echo_addr = app
+                .instantiate_contract(code_id, owner.clone(), &EmptyMsg {}, &[], "label", None)
+                .unwrap();
+
+            // ensure the execute has the same wrapper as it should
+            let msg = echo::Message::<Empty> {
+                data: Some("hello".into()),
+                ..echo::Message::default()
+            };
+            let exec_res = app.execute_contract(owner, echo_addr, &msg, &[]).unwrap();
+            assert!(exec_res.data.is_some());
+            let parsed = parse_execute_response_data(&exec_res.data.unwrap()).unwrap();
+            assert_eq!(parsed.data, Some(Binary::from(b"hello")));
+        }
+    }
 }

packages/multi-test/src/executor.rs

@@ -1,10 +1,10 @@
 use std::fmt;
 
-use crate::parse_contract_addr;
 use cosmwasm_std::{
     to_binary, Addr, Attribute, BankMsg, Binary, Coin, CosmosMsg, Event, SubMsgExecutionResponse,
     WasmMsg,
 };
+use cw0::parse_instantiate_response_data;
 use schemars::JsonSchema;
 use serde::Serialize;
 
@@ -91,7 +91,8 @@ where
             label: label.into(),
         };
         let res = self.execute(sender, msg.into())?;
-        parse_contract_addr(&res.data)
+        let data = parse_instantiate_response_data(res.data.unwrap_or_default().as_slice())?;
+        Ok(Addr::unchecked(data.contract_address))
     }
 
     /// Execute a contract and process all returned messages.

packages/multi-test/src/lib.rs

@@ -28,4 +28,4 @@ pub use crate::contracts::{Contract, ContractWrapper};
 pub use crate::executor::{AppResponse, Executor};
 pub use crate::module::Module;
 pub use crate::staking::{FailingDistribution, FailingStaking, Staking, StakingSudo};
-pub use crate::wasm::{parse_contract_addr, Wasm, WasmKeeper, WasmSudo};
+pub use crate::wasm::{Wasm, WasmKeeper, WasmSudo};

packages/multi-test/src/test_helpers/contracts/echo.rs

@@ -27,17 +27,35 @@ where
     pub events: Vec<Event>,
 }
 
+// This can take some data... but happy to accept {}
+#[derive(Debug, Clone, Serialize, Deserialize, Derivative)]
+#[derivative(Default(bound = "", new = "true"))]
+pub struct InitMessage<ExecC>
+where
+    ExecC: Debug + PartialEq + Clone + JsonSchema + 'static,
+{
+    pub data: Option<String>,
+    pub sub_msg: Option<Vec<SubMsg<ExecC>>>,
+}
+
 #[allow(clippy::unnecessary_wraps)]
 fn instantiate<ExecC>(
     _deps: DepsMut,
     _env: Env,
     _info: MessageInfo,
-    _msg: EmptyMsg,
+    msg: InitMessage<ExecC>,
 ) -> Result<Response<ExecC>, StdError>
 where
     ExecC: Debug + PartialEq + Clone + JsonSchema + 'static,
 {
-    Ok(Response::default())
+    let mut res = Response::new();
+    if let Some(data) = msg.data {
+        res = res.set_data(data.into_bytes());
+    }
+    if let Some(msgs) = msg.sub_msg {
+        res = res.add_submessages(msgs);
+    }
+    Ok(res)
 }
 
 #[allow(clippy::unnecessary_wraps)]

packages/multi-test/src/wasm.rs

@@ -23,8 +23,7 @@ use crate::executor::AppResponse;
 use crate::transactions::transactional;
 use cosmwasm_std::testing::mock_wasmd_attr;
 
-use anyhow::{anyhow, bail, Result as AnyResult};
-use cw0::parse_instantiate_response_data;
+use anyhow::{bail, Result as AnyResult};
 
 // Contract state is kept in Storage, separate from the contracts themselves
 const CONTRACTS: Map<&Addr, ContractData> = Map::new("contracts");
@@ -149,11 +148,7 @@ where
         sender: Addr,
         msg: WasmMsg,
     ) -> AnyResult<AppResponse> {
-        let (resender, res, custom_event) =
-            self.execute_wasm(api, storage, router, block, sender, msg)?;
-
-        let (res, msgs) = self.build_app_response(&resender, custom_event, res);
-        self.process_response(api, router, storage, block, resender, res, msgs)
+        self.execute_wasm(api, storage, router, block, sender, msg)
     }
 
     fn sudo(
@@ -256,7 +251,7 @@ where
         block: &BlockInfo,
         sender: Addr,
         wasm_msg: WasmMsg,
-    ) -> AnyResult<(Addr, Response<ExecC>, Event)> {
+    ) -> AnyResult<AppResponse> {
         match wasm_msg {
             WasmMsg::Execute {
                 contract_addr,
@@ -289,7 +284,12 @@ where
 
                 let custom_event =
                     Event::new("execute").add_attribute(CONTRACT_ATTR, &contract_addr);
-                Ok((contract_addr, res, custom_event))
+
+                let (res, msgs) = self.build_app_response(&contract_addr, custom_event, res);
+                let res =
+                    self.process_response(api, router, storage, block, contract_addr, res, msgs)?;
+                // res.data = execute_response(res.data);
+                Ok(res)
             }
             WasmMsg::Instantiate {
                 admin,
@@ -324,7 +324,7 @@ where
 
                 // then call the contract
                 let info = MessageInfo { sender, funds };
-                let mut res = self.call_instantiate(
+                let res = self.call_instantiate(
                     contract_addr.clone(),
                     api,
                     storage,
@@ -333,12 +333,23 @@ where
                     info,
                     msg.to_vec(),
                 )?;
-                init_response(&mut res, &contract_addr);
 
                 let custom_event = Event::new("instantiate")
                     .add_attribute(CONTRACT_ATTR, &contract_addr)
                     .add_attribute("code_id", code_id.to_string());
-                Ok((contract_addr, res, custom_event))
+
+                let (res, msgs) = self.build_app_response(&contract_addr, custom_event, res);
+                let mut res = self.process_response(
+                    api,
+                    router,
+                    storage,
+                    block,
+                    contract_addr.clone(),
+                    res,
+                    msgs,
+                )?;
+                res.data = Some(init_response(res.data, &contract_addr));
+                Ok(res)
             }
             WasmMsg::Migrate {
                 contract_addr,
@@ -372,7 +383,11 @@ where
                 let custom_event = Event::new("migrate")
                     .add_attribute(CONTRACT_ATTR, &contract_addr)
                     .add_attribute("code_id", new_code_id.to_string());
-                Ok((contract_addr, res, custom_event))
+                let (res, msgs) = self.build_app_response(&contract_addr, custom_event, res);
+                let res =
+                    self.process_response(api, router, storage, block, contract_addr, res, msgs)?;
+                // res.data = execute_response(res.data);
+                Ok(res)
             }
             msg => bail!(Error::UnsupportedWasmMsg(msg)),
         }
@@ -826,42 +841,45 @@ where
     }
 }
 
+// TODO: replace with code in cw0
+
 #[derive(Clone, PartialEq, Message)]
-pub struct InstantiateData {
+struct InstantiateResponse {
     #[prost(string, tag = "1")]
     pub address: ::prost::alloc::string::String,
-    /// Unique ID number for this person.
     #[prost(bytes, tag = "2")]
     pub data: ::prost::alloc::vec::Vec<u8>,
 }
 
-fn init_response<C>(res: &mut Response<C>, contact_address: &Addr)
-where
-    C: Clone + fmt::Debug + PartialEq + JsonSchema,
-{
-    let data = res.data.clone().unwrap_or_default().to_vec();
-    let init_data = InstantiateData {
+// TODO: encode helpers in cw0
+fn init_response(data: Option<Binary>, contact_address: &Addr) -> Binary {
+    let data = data.unwrap_or_default().to_vec();
+    let init_data = InstantiateResponse {
         address: contact_address.into(),
         data,
     };
     let mut new_data = Vec::<u8>::with_capacity(init_data.encoded_len());
     // the data must encode successfully
     init_data.encode(&mut new_data).unwrap();
-    res.data = Some(new_data.into());
+    new_data.into()
 }
 
-// this parses the result from a wasm contract init
-pub fn parse_contract_addr(data: &Option<Binary>) -> AnyResult<Addr> {
-    let bin = data
-        .as_ref()
-        .ok_or_else(|| anyhow!("No data response"))?
-        .to_vec();
-    // parse the protobuf struct
-    let init_data = parse_instantiate_response_data(bin.as_slice())?;
-    if init_data.contract_address.is_empty() {
-        bail!("no contract address provided");
-    }
-    Ok(Addr::unchecked(init_data.contract_address))
+#[derive(Clone, PartialEq, Message)]
+struct ExecuteResponse {
+    #[prost(bytes, tag = "1")]
+    pub data: ::prost::alloc::vec::Vec<u8>,
+}
+
+// empty return if no data present in original
+#[allow(dead_code)]
+fn execute_response(data: Option<Binary>) -> Option<Binary> {
+    data.map(|d| {
+        let exec_data = ExecuteResponse { data: d.to_vec() };
+        let mut new_data = Vec::<u8>::with_capacity(exec_data.encoded_len());
+        // the data must encode successfully
+        exec_data.encode(&mut new_data).unwrap();
+        new_data.into()
+    })
 }
 
 #[cfg(test)]