diff --git a/contracts/voting/dao-voting-cw721-staked/src/contract.rs b/contracts/voting/dao-voting-cw721-staked/src/contract.rs index fc83cce13..df35cd873 100644 --- a/contracts/voting/dao-voting-cw721-staked/src/contract.rs +++ b/contracts/voting/dao-voting-cw721-staked/src/contract.rs @@ -17,7 +17,7 @@ use dao_voting::threshold::{ ActiveThresholdResponse, }; -use crate::msg::{ExecuteMsg, InstantiateMsg, MigrateMsg, NftContract, QueryMsg}; +use crate::msg::{ExecuteMsg, InstantiateMsg, MigrateMsg, NftContract, QueryMsg, AddressResponse}; use crate::state::{ register_staked_nft, register_unstaked_nfts, Config, ACTIVE_THRESHOLD, CONFIG, DAO, HOOKS, INITIAL_NFTS, MAX_CLAIMS, NFT_BALANCES, NFT_CLAIMS, STAKED_NFTS_PER_OWNER, TOTAL_STAKED_NFTS, @@ -495,6 +495,53 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult { QueryMsg::TotalPowerAtHeight { height } => query_total_power_at_height(deps, env, height), QueryMsg::VotingPowerAtHeight { address, height } => { query_voting_power_at_height(deps, env, address, height) + }, + QueryMsg::NftOwner { token_id } => query_nft_owner(deps, token_id), + } +} + +pub fn query_nft_owner(deps: Deps, token_id: String) -> StdResult { + let config = CONFIG.load(deps.storage)?; + let nft_contract = config.nft_address; + + // Querying the NFT contract for the owner of the token + let owner: StdResult = deps.querier.query_wasm_smart( + nft_contract, + &cw721::Cw721QueryMsg::OwnerOf { + token_id: token_id.clone(), + include_expired: None, + }, + ); + + match owner { + Ok(owner_info) => { + // Checking if the token is staked in this contract + let staker = STAKED_NFTS_PER_OWNER + .prefix_range(deps.storage, None, None, cosmwasm_std::Order::Ascending) + .find_map(|result| { + let ((addr, staked_token_id), _) = result.unwrap(); + if staked_token_id == token_id { + Some(addr) + } else { + None + } + }); + + let response = if let Some(staker) = staker { + AddressResponse { + owner: Some(staker.to_string()), + } + } else { + AddressResponse { + owner: Some(owner_info.owner), + } + }; + + to_json_binary(&response) + }, + Err(_) => { + // If the NFT doesn't exist, return None for the owner + to_json_binary(&AddressResponse { owner: None }) } } } diff --git a/contracts/voting/dao-voting-cw721-staked/src/msg.rs b/contracts/voting/dao-voting-cw721-staked/src/msg.rs index 837851ed3..6760b9819 100644 --- a/contracts/voting/dao-voting-cw721-staked/src/msg.rs +++ b/contracts/voting/dao-voting-cw721-staked/src/msg.rs @@ -92,7 +92,14 @@ pub enum QueryMsg { }, #[returns(ActiveThresholdResponse)] ActiveThreshold {}, + #[returns(AddressResponse)] + NftOwner { token_id: String }, } #[cw_serde] pub struct MigrateMsg {} + +#[cw_serde] +pub struct AddressResponse { + pub owner: Option, +} \ No newline at end of file diff --git a/contracts/voting/dao-voting-cw721-staked/src/testing/tests.rs b/contracts/voting/dao-voting-cw721-staked/src/testing/tests.rs index 7126bd904..d2b7c3fff 100644 --- a/contracts/voting/dao-voting-cw721-staked/src/testing/tests.rs +++ b/contracts/voting/dao-voting-cw721-staked/src/testing/tests.rs @@ -21,6 +21,7 @@ use crate::{ queries::{query_config, query_hooks, query_nft_owner, query_total_and_voting_power}, }, }; +use crate::msg::AddressResponse; use super::instantiate::instantiate_cw721_base; use super::{ @@ -1453,3 +1454,103 @@ pub fn test_migrate_update_version() { assert_eq!(version.version, CONTRACT_VERSION); assert_eq!(version.contract, CONTRACT_NAME); } + +#[test] +fn test_query_nft_owner() -> anyhow::Result<()> { + let CommonTest { + mut app, + module, + nft, + } = setup_test(None); + + // Mint and stake an NFT + mint_and_stake_nft(&mut app, &nft, &module, CREATOR_ADDR, "1")?; + + // Query the owner of the staked NFT + let res: AddressResponse = app + .wrap() + .query_wasm_smart(module.clone(), &QueryMsg::NftOwner { token_id: "1".to_string() })?; + + assert_eq!(res.owner, Some(CREATOR_ADDR.to_string())); + + // Mint an NFT but don't stake it + mint_nft(&mut app, &nft, CREATOR_ADDR, "recipient", "2")?; + + // Query the owner of the unstaked NFT + let res: AddressResponse = app + .wrap() + .query_wasm_smart(module.clone(), &QueryMsg::NftOwner { token_id: "2".to_string() })?; + + assert_eq!(res.owner, Some("recipient".to_string())); + + // Query a non-existent NFT + let res: AddressResponse = app + .wrap() + .query_wasm_smart(module, &QueryMsg::NftOwner { token_id: "999".to_string() })?; + + assert_eq!(res.owner, None, "Querying non-existent NFT should return None for owner"); + + Ok(()) +} + +#[test] +fn test_query_nft_owner_after_unstake() -> anyhow::Result<()> { + let CommonTest { + mut app, + module, + nft, + } = setup_test(None); + + // Mint and stake an NFT + mint_and_stake_nft(&mut app, &nft, &module, CREATOR_ADDR, "1")?; + + + // Query the owner of the staked NFT. + let res: AddressResponse = app + .wrap() + .query_wasm_smart(module.clone(), &QueryMsg::NftOwner { token_id: "1".to_string() })?; + + + assert_eq!(res.owner, Some(CREATOR_ADDR.to_string())); + + // Unstake the NFT + unstake_nfts(&mut app, &module, CREATOR_ADDR, &["1"])?; + + // Query the owner after unstaking + let res: AddressResponse = app + .wrap() + .query_wasm_smart(module, &QueryMsg::NftOwner { token_id: "1".to_string() })?; + + assert_eq!(res.owner, Some(CREATOR_ADDR.to_string())); + + Ok(()) +} + +#[test] +fn test_query_nft_owner_with_multiple_stakers() -> anyhow::Result<()> { + let CommonTest { + mut app, + module, + nft, + } = setup_test(None); + + // Mint and stake NFTs for different users. + mint_and_stake_nft(&mut app, &nft, &module, CREATOR_ADDR, "1")?; + mint_nft(&mut app, &nft, CREATOR_ADDR, "user2", "2")?; + stake_nft(&mut app, &nft, &module, "user2", "2")?; + + // Query owners. + let res1: AddressResponse = app + .wrap() + .query_wasm_smart(module.clone(), &QueryMsg::NftOwner { token_id: "1".to_string() })?; + + assert_eq!(res1.owner, Some(CREATOR_ADDR.to_string())); + + let res2: AddressResponse = app + .wrap() + .query_wasm_smart(module, &QueryMsg::NftOwner { token_id: "2".to_string() })?; + + assert_eq!(res2.owner, Some("user2".to_string())); + + Ok(()) +} \ No newline at end of file