Refactor query logic

contracts/ics721/src/contract.rs

@@ -1,18 +1,14 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
-use cosmwasm_std::{
-    from_binary, to_binary, Addr, Binary, Deps, DepsMut, Env, IbcMsg, IbcQuery, MessageInfo, Order,
-    PortIdResponse, Response, StdResult,
-};
+use cosmwasm_std::{from_binary, to_binary, Addr, DepsMut, Env, IbcMsg, MessageInfo, Response};
 use cw2::set_contract_version;
-use cw20_ics20::msg::{ListChannelsResponse, PortResponse};
 use cw721_ibc::Cw721ReceiveMsg;
 use cw_utils::nonpayable;
 
 use crate::error::ContractError;
 use crate::ibc::Ics721Packet;
-use crate::msg::{ChannelResponse, ExecuteMsg, InstantiateMsg, QueryMsg, TransferMsg};
-use crate::state::{Config, CHANNEL_INFO, CHANNEL_STATE, CONFIG};
+use crate::msg::{ExecuteMsg, InstantiateMsg, TransferMsg};
+use crate::state::{Config, CHANNEL_INFO, CONFIG};
 
 // version info for migration info
 const CONTRACT_NAME: &str = "crates.io:sg721-ics721";
@@ -118,68 +114,3 @@ pub fn execute_transfer(
         .add_attribute("token_ids", &packet.token_ids.join(","));
     Ok(res)
 }
-
-// TODO: Alot of this query code is copy pasta.
-// Find a way to make it generic or put into a package.
-#[cfg_attr(not(feature = "library"), entry_point)]
-pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
-    match msg {
-        QueryMsg::Port {} => to_binary(&query_port(deps)?),
-        QueryMsg::ListChannels {} => to_binary(&query_list(deps)?),
-        QueryMsg::Channel { id } => to_binary(&query_channel(deps, id)?),
-        QueryMsg::Tokens {
-            channel_id,
-            class_id,
-        } => to_binary(&query_tokens(deps, channel_id, class_id)?),
-    }
-}
-
-fn query_port(deps: Deps) -> StdResult<PortResponse> {
-    let query = IbcQuery::PortId {}.into();
-    let PortIdResponse { port_id } = deps.querier.query(&query)?;
-    Ok(PortResponse { port_id })
-}
-
-fn query_list(deps: Deps) -> StdResult<ListChannelsResponse> {
-    let channels: StdResult<Vec<_>> = CHANNEL_INFO
-        .range(deps.storage, None, None, Order::Ascending)
-        .map(|r| r.map(|(_, v)| v))
-        .collect();
-    Ok(ListChannelsResponse {
-        channels: channels?,
-    })
-}
-
-pub fn query_channel(deps: Deps, id: String) -> StdResult<ChannelResponse> {
-    let info = CHANNEL_INFO.load(deps.storage, &id)?;
-    let _class_ids: StdResult<Vec<_>> = CHANNEL_STATE
-        .sub_prefix(&id)
-        .range(deps.storage, None, None, Order::Ascending)
-        .map(|r| {
-            let (class_id_token_id, _) = r?;
-            Ok(class_id_token_id.0)
-        })
-        .collect();
-
-    let class_ids_resp = _class_ids;
-    match class_ids_resp {
-        Ok(mut class_id_vec) => Ok(ChannelResponse {
-            info,
-            class_ids: {
-                class_id_vec.sort();
-                class_id_vec.dedup();
-                class_id_vec
-            },
-        }),
-        Err(msg) => Err(msg),
-    }
-}
-
-// TODO: https://github.com/public-awesome/contracts/issues/59
-pub fn query_tokens(
-    _deps: Deps,
-    _channel_id: String,
-    _class_id: String,
-) -> StdResult<ChannelResponse> {
-    todo!()
-}

contracts/ics721/src/contract_test.rs

@@ -1,15 +1,26 @@
 #[cfg(test)]
 mod contact_testing {
+    use core::panic;
+
     use super::super::*;
+    use crate::msg::{ChannelResponse, ExecuteMsg, InstantiateMsg, QueryMsg, TransferMsg};
+    use crate::query;
+    use crate::state::CHANNEL_STATE;
     use crate::test_constants::*;
     use crate::test_helpers::*;
 
     use cosmwasm_std::testing::mock_env;
-    use cosmwasm_std::{from_binary, to_binary, Attribute, Coin, StdError};
+    use cosmwasm_std::Deps;
+    use cosmwasm_std::{
+        from_binary, to_binary, Addr, Attribute, Coin, MessageInfo, Response, StdError,
+    };
     use cosmwasm_std::{CosmosMsg, IbcEndpoint};
     use cw2::{get_contract_version, ContractVersion};
+    use cw20_ics20::msg::ListChannelsResponse;
     use cw20_ics20::state::ChannelInfo;
+    use cw20_ics20::ContractError;
     use cw721_ibc::Cw721ReceiveMsg;
+    use query::{query, query_channel};
 
     use cosmwasm_std::testing::mock_dependencies;
 
@@ -59,9 +70,8 @@ mod contact_testing {
     #[test]
     fn test_query_channel_list_success() {
         let deps = setup(&[TEST_CHANNEL_0_DATA, TEST_CHANNEL_1_DATA]);
-        let result = query_list(deps.as_ref());
-
-        let expected_list: StdResult<ListChannelsResponse> = Ok(ListChannelsResponse {
+        let result = query(deps.as_ref(), mock_env(), QueryMsg::ListChannels {});
+        let expected_list = ListChannelsResponse {
             channels: vec![
                 ChannelInfo {
                     id: CHANNEL_FROM_STARS_TO_OMNI.to_string(),
@@ -80,20 +90,27 @@ mod contact_testing {
                     connection_id: CONNECTION_1.to_string(),
                 },
             ],
-        });
-        assert_eq!(result, expected_list);
+        };
+        let expected_list_bin = to_binary(&expected_list).unwrap();
+        match result {
+            Ok(bin) => assert_eq!(bin, expected_list_bin),
+            Err(_err) => panic!("Query failed for test_query_channel_list_success"),
+        }
     }
 
     #[test]
     fn test_query_channel_list_empty() {
         let mut deps = setup(&[TEST_CHANNEL_0_DATA, TEST_CHANNEL_1_DATA]);
         CHANNEL_INFO.remove(&mut deps.storage, CHANNEL_FROM_STARS_TO_OMNI);
         CHANNEL_INFO.remove(&mut deps.storage, CHANNEL_FROM_STARS_TO_GB);
-        let result = query_list(deps.as_ref());
+        let result = query(deps.as_ref(), mock_env(), QueryMsg::ListChannels {});
 
-        let expected_list: StdResult<ListChannelsResponse> =
-            Ok(ListChannelsResponse { channels: vec![] });
-        assert_eq!(result, expected_list);
+        let expected_list = ListChannelsResponse { channels: vec![] };
+        let expected_list_bin = to_binary(&expected_list).unwrap();
+        match result {
+            Ok(bin) => assert_eq!(bin, expected_list_bin),
+            Err(_err) => panic!("Query failed for test_query_channel_list_empty"),
+        }
     }
 
     #[test]

contracts/ics721/src/ibc_test.rs

@@ -4,15 +4,17 @@ mod ibc_testing {
     use std::vec;
 
     use super::super::*;
+    use crate::query;
     use crate::test_constants::{
         CHANNEL_FROM_OMNI_TO_STARS, CHANNEL_FROM_STARS_TO_OMNI, CONNECTION_0, TEST_CHANNEL_0_DATA,
         TEST_CHANNEL_1_DATA,
     };
     use crate::test_helpers::*;
+
     use cosmwasm_std::CosmosMsg::Wasm;
     use cosmwasm_std::WasmMsg::Execute;
+    use query::query_channel;
 
-    use crate::contract::query_channel;
     use cosmwasm_std::testing::mock_dependencies;
     use cosmwasm_std::testing::mock_env;
     use cosmwasm_std::{

contracts/ics721/src/lib.rs

@@ -2,6 +2,7 @@ pub mod contract;
 mod error;
 pub mod ibc;
 pub mod msg;
+pub mod query;
 pub mod state;
 mod test_constants;
 mod test_helpers;

contracts/ics721/src/query.rs

@@ -0,0 +1,70 @@
+use crate::msg::{ChannelResponse, QueryMsg};
+use crate::state::{CHANNEL_INFO, CHANNEL_STATE};
+
+#[cfg(not(feature = "library"))]
+use cosmwasm_std::entry_point;
+use cosmwasm_std::{to_binary, Binary, Deps, Env, IbcQuery, Order, PortIdResponse, StdResult};
+use cw20_ics20::msg::{ListChannelsResponse, PortResponse};
+
+#[cfg_attr(not(feature = "library"), entry_point)]
+pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
+    match msg {
+        QueryMsg::Port {} => to_binary(&query_port(deps)?),
+        QueryMsg::ListChannels {} => to_binary(&query_list(deps)?),
+        QueryMsg::Channel { id } => to_binary(&query_channel(deps, id)?),
+        QueryMsg::Tokens {
+            channel_id,
+            class_id,
+        } => to_binary(&query_tokens(deps, channel_id, class_id)?),
+    }
+}
+
+fn query_port(deps: Deps) -> StdResult<PortResponse> {
+    let query = IbcQuery::PortId {}.into();
+    let PortIdResponse { port_id } = deps.querier.query(&query)?;
+    Ok(PortResponse { port_id })
+}
+
+fn query_list(deps: Deps) -> StdResult<ListChannelsResponse> {
+    let channels: StdResult<Vec<_>> = CHANNEL_INFO
+        .range(deps.storage, None, None, Order::Ascending)
+        .map(|r| r.map(|(_, v)| v))
+        .collect();
+    Ok(ListChannelsResponse {
+        channels: channels?,
+    })
+}
+
+pub fn query_channel(deps: Deps, id: String) -> StdResult<ChannelResponse> {
+    let info = CHANNEL_INFO.load(deps.storage, &id)?;
+    let _class_ids: StdResult<Vec<_>> = CHANNEL_STATE
+        .sub_prefix(&id)
+        .range(deps.storage, None, None, Order::Ascending)
+        .map(|r| {
+            let (class_id_token_id, _) = r?;
+            Ok(class_id_token_id.0)
+        })
+        .collect();
+
+    let class_ids_resp = _class_ids;
+    match class_ids_resp {
+        Ok(mut class_id_vec) => Ok(ChannelResponse {
+            info,
+            class_ids: {
+                class_id_vec.sort();
+                class_id_vec.dedup();
+                class_id_vec
+            },
+        }),
+        Err(msg) => Err(msg),
+    }
+}
+
+// TODO: https://github.com/public-awesome/contracts/issues/59
+pub fn query_tokens(
+    _deps: Deps,
+    _channel_id: String,
+    _class_id: String,
+) -> StdResult<ChannelResponse> {
+    todo!()
+}