1use cosmwasm_std::Addr;
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4
5#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
7pub struct Config {
8 pub owner: Addr,
10 pub council_address: Addr,
12 pub incentives_address: Addr,
14 pub safety_fund_address: Addr,
16 pub mars_token_address: Addr,
18 pub oracle_address: Addr,
20 pub protocol_admin_address: Addr,
22 pub protocol_rewards_collector_address: Addr,
24 pub red_bank_address: Addr,
26 pub staking_address: Addr,
28 pub treasury_address: Addr,
30 pub vesting_address: Addr,
32 pub xmars_token_address: Addr,
34}
35
36#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
37pub enum MarsContract {
39 Council,
40 Incentives,
41 SafetyFund,
42 MarsToken,
43 Oracle,
44 ProtocolAdmin,
45 ProtocolRewardsCollector,
46 RedBank,
47 Staking,
48 Treasury,
49 Vesting,
50 XMarsToken,
51}
52
53pub mod msg {
54 use schemars::JsonSchema;
55 use serde::{Deserialize, Serialize};
56
57 use super::MarsContract;
58
59 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
65 pub struct InstantiateMsg {
66 pub owner: String,
67 }
68
69 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
70 #[serde(rename_all = "snake_case")]
71 pub enum ExecuteMsg {
72 UpdateConfig { config: ConfigParams },
74 }
75
76 #[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq, JsonSchema)]
77 pub struct ConfigParams {
78 pub owner: Option<String>,
80 pub council_address: Option<String>,
82 pub incentives_address: Option<String>,
84 pub safety_fund_address: Option<String>,
87 pub mars_token_address: Option<String>,
89 pub oracle_address: Option<String>,
91 pub protocol_admin_address: Option<String>,
94 pub protocol_rewards_collector_address: Option<String>,
96 pub red_bank_address: Option<String>,
99 pub staking_address: Option<String>,
101 pub treasury_address: Option<String>,
104 pub vesting_address: Option<String>,
106 pub xmars_token_address: Option<String>,
108 }
109
110 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
111 #[serde(rename_all = "snake_case")]
112 pub enum QueryMsg {
113 Config {},
115 Address { contract: MarsContract },
117 Addresses { contracts: Vec<MarsContract> },
119 }
120}
121
122pub mod helpers {
123 use super::msg::QueryMsg;
124 use super::MarsContract;
125 use crate::error::MarsError;
126 use cosmwasm_std::{to_binary, Addr, QuerierWrapper, QueryRequest, WasmQuery};
127
128 pub fn query_address(
129 querier: &QuerierWrapper,
130 address_provider_address: Addr,
131 contract: MarsContract,
132 ) -> Result<Addr, MarsError> {
133 let query: Addr = querier.query(&QueryRequest::Wasm(WasmQuery::Smart {
134 contract_addr: address_provider_address.to_string(),
135 msg: to_binary(&QueryMsg::Address {
136 contract: contract.clone(),
137 })?,
138 }))?;
139
140 if query == Addr::unchecked("") {
141 Err(MarsError::EmptyAddresses {
142 empty_addresses: vec![contract],
143 })
144 } else {
145 Ok(query)
146 }
147 }
148
149 pub fn query_addresses(
150 querier: &QuerierWrapper,
151 address_provider_address: Addr,
152 contracts: Vec<MarsContract>,
153 ) -> Result<Vec<Addr>, MarsError> {
154 let expected_len = contracts.len();
155
156 let query: Vec<Addr> = querier.query(&QueryRequest::Wasm(WasmQuery::Smart {
157 contract_addr: address_provider_address.to_string(),
158 msg: to_binary(&QueryMsg::Addresses {
159 contracts: contracts.clone(),
160 })?,
161 }))?;
162
163 if query.len() != expected_len {
164 return Err(MarsError::AddressesQueryWrongNumber {
165 expected: expected_len as u32,
166 actual: query.len() as u32,
167 });
168 }
169
170 let empty_addresses = query
171 .iter()
172 .zip(contracts)
173 .filter(|(address, _)| *address == &Addr::unchecked(""))
174 .map(|(_, contract)| contract)
175 .collect::<Vec<MarsContract>>();
176
177 if !empty_addresses.is_empty() {
178 Err(MarsError::EmptyAddresses { empty_addresses })
179 } else {
180 Ok(query)
181 }
182 }
183}
184
185#[cfg(test)]
188mod tests {
189 use super::*;
190 use crate::address_provider::msg::QueryMsg;
191 use crate::error::MarsError;
192 use cosmwasm_std::testing::{MockApi, MockStorage};
193 use cosmwasm_std::{
194 from_binary, from_slice, to_binary, Binary, ContractResult, OwnedDeps, Querier,
195 QuerierResult, QueryRequest, StdResult, SystemError, WasmQuery,
196 };
197 use terra_cosmwasm::TerraQueryWrapper;
198
199 #[test]
200 fn test_query_address() {
201 let deps = OwnedDeps {
202 storage: MockStorage::default(),
203 api: MockApi::default(),
204 querier: AddressProviderMockQuerier {},
205 };
206
207 {
209 let err = helpers::query_address(
210 &deps.as_ref().querier,
211 Addr::unchecked("address_provider"),
212 MarsContract::Incentives,
213 )
214 .unwrap_err();
215
216 assert_eq!(
217 err,
218 MarsError::EmptyAddresses {
219 empty_addresses: vec![MarsContract::Incentives]
220 }
221 );
222 }
223
224 {
226 let address = helpers::query_address(
227 &deps.as_ref().querier,
228 Addr::unchecked("address_provider"),
229 MarsContract::RedBank,
230 )
231 .unwrap();
232
233 assert_eq!(address, Addr::unchecked("red_bank"));
234 }
235 }
236
237 #[test]
238 fn test_query_addresses() {
239 let deps = OwnedDeps {
240 storage: MockStorage::default(),
241 api: MockApi::default(),
242 querier: AddressProviderMockQuerier {},
243 };
244
245 {
247 let err = helpers::query_addresses(
248 &deps.as_ref().querier,
249 Addr::unchecked("address_provider"),
250 vec![
251 MarsContract::ProtocolRewardsCollector,
252 MarsContract::RedBank,
253 MarsContract::Incentives,
254 ],
255 )
256 .unwrap_err();
257
258 assert_eq!(
259 err,
260 MarsError::EmptyAddresses {
261 empty_addresses: vec![
262 MarsContract::ProtocolRewardsCollector,
263 MarsContract::Incentives
264 ]
265 }
266 );
267 }
268
269 {
271 let addresses = helpers::query_addresses(
272 &deps.as_ref().querier,
273 Addr::unchecked("address_provider"),
274 vec![MarsContract::Vesting, MarsContract::RedBank],
275 )
276 .unwrap();
277
278 assert_eq!(
279 addresses,
280 vec![Addr::unchecked("vesting"), Addr::unchecked("red_bank")]
281 );
282 }
283 }
284
285 #[derive(Clone, Copy)]
286 pub struct AddressProviderMockQuerier {}
287
288 impl Querier for AddressProviderMockQuerier {
289 fn raw_query(&self, bin_request: &[u8]) -> QuerierResult {
290 let request: QueryRequest<TerraQueryWrapper> = match from_slice(bin_request) {
291 Ok(v) => v,
292 Err(e) => {
293 return Err(SystemError::InvalidRequest {
294 error: format!("Parsing query request: {}", e),
295 request: bin_request.into(),
296 })
297 .into()
298 }
299 };
300 self.handle_query(&request)
301 }
302 }
303
304 impl AddressProviderMockQuerier {
305 pub fn handle_query(&self, request: &QueryRequest<TerraQueryWrapper>) -> QuerierResult {
306 if let QueryRequest::Wasm(WasmQuery::Smart {
307 contract_addr: _,
308 msg,
309 }) = request
310 {
311 let parse_address_provider_query: StdResult<QueryMsg> = from_binary(msg);
312
313 if let Ok(address_provider_query) = parse_address_provider_query {
314 let ret: ContractResult<Binary> = match address_provider_query {
315 QueryMsg::Address { contract } => {
316 to_binary(&get_contract_address(contract)).into()
317 }
318
319 QueryMsg::Addresses { contracts } => {
320 let addresses = contracts
321 .into_iter()
322 .map(get_contract_address)
323 .collect::<Vec<_>>();
324 to_binary(&addresses).into()
325 }
326
327 _ => panic!("[mock]: Unsupported address provider query"),
328 };
329
330 return Ok(ret).into();
331 }
332 }
333
334 panic!("[mock]: Unsupported wasm query");
335 }
336 }
337
338 fn get_contract_address(contract: MarsContract) -> Addr {
339 match contract {
340 MarsContract::Incentives => Addr::unchecked(""),
342 MarsContract::ProtocolRewardsCollector => Addr::unchecked(""),
343
344 MarsContract::Council => Addr::unchecked("council"),
346 MarsContract::SafetyFund => Addr::unchecked("safety_fund"),
347 MarsContract::MarsToken => Addr::unchecked("mars_token"),
348 MarsContract::Oracle => Addr::unchecked("oracle"),
349 MarsContract::ProtocolAdmin => Addr::unchecked("protocol_admin"),
350 MarsContract::RedBank => Addr::unchecked("red_bank"),
351 MarsContract::Staking => Addr::unchecked("staking"),
352 MarsContract::Treasury => Addr::unchecked("treasury"),
353 MarsContract::Vesting => Addr::unchecked("vesting"),
354 MarsContract::XMarsToken => Addr::unchecked("xmars_token"),
355 }
356 }
357}