cw_token_swap/
state.rs

1use cosmwasm_schema::cw_serde;
2use cosmwasm_std::{
3    to_json_binary, Addr, BankMsg, Coin, CosmosMsg, Deps, StdError, Uint128, WasmMsg,
4};
5use cw_storage_plus::Item;
6
7use crate::{
8    msg::{Counterparty, TokenInfo},
9    ContractError,
10};
11
12#[cw_serde]
13pub enum CheckedTokenInfo {
14    Native {
15        denom: String,
16        amount: Uint128,
17    },
18    Cw20 {
19        contract_addr: Addr,
20        amount: Uint128,
21    },
22}
23
24#[cw_serde]
25pub struct CheckedCounterparty {
26    pub address: Addr,
27    pub promise: CheckedTokenInfo,
28    pub provided: bool,
29}
30
31pub const COUNTERPARTY_ONE: Item<CheckedCounterparty> = Item::new("counterparty_one");
32pub const COUNTERPARTY_TWO: Item<CheckedCounterparty> = Item::new("counterparty_two");
33
34impl Counterparty {
35    pub fn into_checked(self, deps: Deps) -> Result<CheckedCounterparty, ContractError> {
36        Ok(CheckedCounterparty {
37            address: deps.api.addr_validate(&self.address)?,
38            provided: false,
39            promise: self.promise.into_checked(deps)?,
40        })
41    }
42}
43
44impl TokenInfo {
45    pub fn into_checked(self, deps: Deps) -> Result<CheckedTokenInfo, ContractError> {
46        match self {
47            TokenInfo::Native { denom, amount } => {
48                if amount.is_zero() {
49                    Err(ContractError::ZeroTokens {})
50                } else {
51                    Ok(CheckedTokenInfo::Native { denom, amount })
52                }
53            }
54            TokenInfo::Cw20 {
55                contract_addr,
56                amount,
57            } => {
58                if amount.is_zero() {
59                    Err(ContractError::ZeroTokens {})
60                } else {
61                    let contract_addr = deps.api.addr_validate(&contract_addr)?;
62                    // Make sure we are dealing with a cw20.
63                    let _: cw20::TokenInfoResponse = deps.querier.query_wasm_smart(
64                        contract_addr.clone(),
65                        &cw20::Cw20QueryMsg::TokenInfo {},
66                    )?;
67                    Ok(CheckedTokenInfo::Cw20 {
68                        contract_addr,
69                        amount,
70                    })
71                }
72            }
73        }
74    }
75}
76
77impl CheckedTokenInfo {
78    pub fn into_send_message(self, recipient: &Addr) -> Result<CosmosMsg, StdError> {
79        Ok(match self {
80            Self::Native { denom, amount } => BankMsg::Send {
81                to_address: recipient.to_string(),
82                amount: vec![Coin { denom, amount }],
83            }
84            .into(),
85            Self::Cw20 {
86                contract_addr,
87                amount,
88            } => WasmMsg::Execute {
89                contract_addr: contract_addr.into_string(),
90                msg: to_json_binary(&cw20::Cw20ExecuteMsg::Transfer {
91                    recipient: recipient.to_string(),
92                    amount,
93                })?,
94                funds: vec![],
95            }
96            .into(),
97        })
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_into_spend_message_native() {
107        let info = CheckedTokenInfo::Native {
108            amount: Uint128::new(100),
109            denom: "uekez".to_string(),
110        };
111        let message = info.into_send_message(&Addr::unchecked("ekez")).unwrap();
112
113        assert_eq!(
114            message,
115            CosmosMsg::Bank(BankMsg::Send {
116                to_address: "ekez".to_string(),
117                amount: vec![Coin {
118                    amount: Uint128::new(100),
119                    denom: "uekez".to_string()
120                }]
121            })
122        );
123    }
124
125    #[test]
126    fn test_into_spend_message_cw20() {
127        let info = CheckedTokenInfo::Cw20 {
128            amount: Uint128::new(100),
129            contract_addr: Addr::unchecked("ekez_token"),
130        };
131        let message = info.into_send_message(&Addr::unchecked("ekez")).unwrap();
132
133        assert_eq!(
134            message,
135            CosmosMsg::Wasm(WasmMsg::Execute {
136                funds: vec![],
137                contract_addr: "ekez_token".to_string(),
138                msg: to_json_binary(&cw20::Cw20ExecuteMsg::Transfer {
139                    recipient: "ekez".to_string(),
140                    amount: Uint128::new(100)
141                })
142                .unwrap()
143            })
144        );
145    }
146}