From 7999ce878eaa16429bba4d1bd344fad1fb770d00 Mon Sep 17 00:00:00 2001
From: Kayanski <kowalski.kowalskin@gmail.com>
Date: Fri, 17 May 2024 08:52:54 +0000
Subject: [PATCH] Added Queryfns on voting

---
 Cargo.lock                                        |  1 +
 contracts/dao-dao-core/src/cw_orch_tests.rs       | 14 ++++----------
 contracts/test/dao-voting-cw20-balance/Cargo.toml |  1 +
 contracts/test/dao-voting-cw20-balance/src/msg.rs |  2 +-
 4 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4a7a7648d..d563d2fe9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2785,6 +2785,7 @@ dependencies = [
  "cosmwasm-schema",
  "cosmwasm-std",
  "cw-multi-test",
+ "cw-orch",
  "cw-storage-plus 1.2.0",
  "cw-utils 1.0.3",
  "cw2 1.1.2",
diff --git a/contracts/dao-dao-core/src/cw_orch_tests.rs b/contracts/dao-dao-core/src/cw_orch_tests.rs
index 66096b85c..aab42583e 100644
--- a/contracts/dao-dao-core/src/cw_orch_tests.rs
+++ b/contracts/dao-dao-core/src/cw_orch_tests.rs
@@ -32,6 +32,7 @@ use dao_interface::{
     voting::{InfoResponse, VotingPowerAtHeightResponse},
 };
 use dao_proposal_sudo::msg::ExecuteMsgFns as _;
+use dao_voting_cw20_balance::msg::QueryMsgFns;
 
 pub fn assert_contains(e: impl std::fmt::Debug, el: impl ToString) {
     assert!(format!("{:?}", e).contains(&el.to_string()))
@@ -1448,11 +1449,8 @@ fn test_cw20_receive_auto_add() {
     voting.set_address(&gov.voting_module().unwrap());
 
     let gov_token = Cw20Base::new("cw20", mock.clone());
-    gov_token.set_address(
-        &voting
-            .query(&dao_voting_cw20_balance::msg::QueryMsg::TokenContract {})
-            .unwrap(),
-    );
+
+    gov_token.set_address(&voting.token_contract().unwrap());
     // Check that the balances query works with no tokens.
     let cw20_balances = gov.cw_20_balances(None, None).unwrap();
     assert_eq!(cw20_balances, vec![]);
@@ -1534,11 +1532,7 @@ fn test_cw20_receive_no_auto_add() {
     voting.set_address(&gov.voting_module().unwrap());
 
     let gov_token = Cw20Base::new("cw20", mock.clone());
-    gov_token.set_address(
-        &voting
-            .query(&dao_voting_cw20_balance::msg::QueryMsg::TokenContract {})
-            .unwrap(),
-    );
+    gov_token.set_address(&voting.token_contract().unwrap());
 
     // Send a gov token to the governance contract. Should not be
     // added becasue auto add is turned off.
diff --git a/contracts/test/dao-voting-cw20-balance/Cargo.toml b/contracts/test/dao-voting-cw20-balance/Cargo.toml
index f126b6840..77e55bf88 100644
--- a/contracts/test/dao-voting-cw20-balance/Cargo.toml
+++ b/contracts/test/dao-voting-cw20-balance/Cargo.toml
@@ -27,6 +27,7 @@ thiserror = { workspace = true }
 dao-dao-macros = { workspace = true }
 dao-interface = { workspace = true }
 cw20-base = {  workspace = true, features = ["library"] }
+cw-orch = "0.22.2"
 
 [dev-dependencies]
 cw-multi-test = { workspace = true }
diff --git a/contracts/test/dao-voting-cw20-balance/src/msg.rs b/contracts/test/dao-voting-cw20-balance/src/msg.rs
index c41590ee2..43002a7ba 100644
--- a/contracts/test/dao-voting-cw20-balance/src/msg.rs
+++ b/contracts/test/dao-voting-cw20-balance/src/msg.rs
@@ -33,5 +33,5 @@ pub enum ExecuteMsg {}
 #[cw20_token_query]
 #[voting_module_query]
 #[cw_serde]
-#[derive(QueryResponses)]
+#[derive(QueryResponses, cw_orch::QueryFns)]
 pub enum QueryMsg {}