|
| 1 | +use { |
| 2 | + crate::{ |
| 3 | + error::PythContractError, |
| 4 | + PriceFeed, |
| 5 | + PriceFeedResponse, |
| 6 | + PriceIdentifier, |
| 7 | + QueryMsg, |
| 8 | + }, |
| 9 | + cosmwasm_std::{ |
| 10 | + from_binary, |
| 11 | + to_binary, |
| 12 | + Binary, |
| 13 | + Coin, |
| 14 | + ContractResult, |
| 15 | + QuerierResult, |
| 16 | + SystemError, |
| 17 | + SystemResult, |
| 18 | + }, |
| 19 | + std::{ |
| 20 | + collections::HashMap, |
| 21 | + time::Duration, |
| 22 | + u128, |
| 23 | + }, |
| 24 | +}; |
| 25 | + |
| 26 | +/// Mock version of Pyth for testing cosmwasm contracts. |
| 27 | +/// This mock stores some price feeds and responds to query messages. |
| 28 | +#[derive(Clone)] |
| 29 | +pub struct MockPyth { |
| 30 | + pub valid_time_period: Duration, |
| 31 | + pub fee_per_vaa: Coin, |
| 32 | + pub feeds: HashMap<PriceIdentifier, PriceFeed>, |
| 33 | +} |
| 34 | + |
| 35 | +impl MockPyth { |
| 36 | + /// Create a new `MockPyth`. You can either provide the full list of price feeds up front, |
| 37 | + /// or add them later via `add_feed`. |
| 38 | + pub fn new(valid_time_period: Duration, fee_per_vaa: Coin, feeds: &[PriceFeed]) -> Self { |
| 39 | + let mut feeds_map = HashMap::new(); |
| 40 | + for feed in feeds { |
| 41 | + feeds_map.insert(feed.id, *feed); |
| 42 | + } |
| 43 | + |
| 44 | + MockPyth { |
| 45 | + valid_time_period, |
| 46 | + fee_per_vaa, |
| 47 | + feeds: feeds_map, |
| 48 | + } |
| 49 | + } |
| 50 | + |
| 51 | + /// Add a price feed that will be returned on queries. |
| 52 | + pub fn add_feed(&mut self, feed: PriceFeed) { |
| 53 | + self.feeds.insert(feed.id, feed); |
| 54 | + } |
| 55 | + |
| 56 | + /// TODO: Update this with example contracts -> Handler for processing query messages. See the tests in `contract.rs` for how to use this |
| 57 | + /// handler within your tests. |
| 58 | + pub fn handle_wasm_query(&self, msg: &Binary) -> QuerierResult { |
| 59 | + let query_msg = from_binary::<QueryMsg>(msg); |
| 60 | + match query_msg { |
| 61 | + Ok(QueryMsg::PriceFeed { id }) => match self.feeds.get(&id) { |
| 62 | + Some(feed) => { |
| 63 | + SystemResult::Ok(to_binary(&PriceFeedResponse { price_feed: *feed }).into()) |
| 64 | + } |
| 65 | + None => SystemResult::Ok(ContractResult::from(Err( |
| 66 | + PythContractError::PriceFeedNotFound, |
| 67 | + ))), |
| 68 | + }, |
| 69 | + Ok(QueryMsg::GetValidTimePeriod) => { |
| 70 | + SystemResult::Ok(to_binary(&self.valid_time_period).into()) |
| 71 | + } |
| 72 | + Ok(QueryMsg::GetUpdateFee { vaas }) => { |
| 73 | + let new_amount = self |
| 74 | + .fee_per_vaa |
| 75 | + .amount |
| 76 | + .u128() |
| 77 | + .checked_mul(vaas.len() as u128) |
| 78 | + .unwrap(); |
| 79 | + SystemResult::Ok(to_binary(&Coin::new(new_amount, &self.fee_per_vaa.denom)).into()) |
| 80 | + } |
| 81 | + Err(_e) => SystemResult::Err(SystemError::InvalidRequest { |
| 82 | + error: "Invalid message".into(), |
| 83 | + request: msg.clone(), |
| 84 | + }), |
| 85 | + } |
| 86 | + } |
| 87 | +} |
0 commit comments