pyth_sdk_cw/
testing.rs

1use {
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)]
29pub struct MockPyth {
30    pub valid_time_period: Duration,
31    pub fee_per_vaa:       Coin,
32    pub feeds:             HashMap<PriceIdentifier, PriceFeed>,
33}
34
35impl 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    /// Handler for processing query messages.
57    /// See the tests in `contract.rs`
58    /// `https://github.com/pyth-network/pyth-crosschain/blob/main/target_chains/cosmwasm/examples/cw-contract/src/contract.rs#L13`
59    /// for how to use this handler within your tests.
60    pub fn handle_wasm_query(&self, msg: &Binary) -> QuerierResult {
61        let query_msg = from_binary::<QueryMsg>(msg);
62        match query_msg {
63            Ok(QueryMsg::PriceFeed { id }) => match self.feeds.get(&id) {
64                Some(feed) => {
65                    SystemResult::Ok(to_binary(&PriceFeedResponse { price_feed: *feed }).into())
66                }
67                None => SystemResult::Ok(ContractResult::from(Err(
68                    PythContractError::PriceFeedNotFound,
69                ))),
70            },
71            Ok(QueryMsg::GetValidTimePeriod) => {
72                SystemResult::Ok(to_binary(&self.valid_time_period).into())
73            }
74
75            Ok(QueryMsg::GetUpdateFee { vaas }) => {
76                let new_amount = self
77                    .fee_per_vaa
78                    .amount
79                    .u128()
80                    .checked_mul(vaas.len() as u128)
81                    .unwrap();
82                SystemResult::Ok(to_binary(&Coin::new(new_amount, &self.fee_per_vaa.denom)).into())
83            }
84            #[cfg(feature = "osmosis")]
85            Ok(QueryMsg::GetUpdateFeeForDenom { vaas, denom }) => {
86                let new_amount = self
87                    .fee_per_vaa
88                    .amount
89                    .u128()
90                    .checked_mul(vaas.len() as u128)
91                    .unwrap();
92                SystemResult::Ok(to_binary(&Coin::new(new_amount, denom)).into())
93            }
94            Err(_e) => SystemResult::Err(SystemError::InvalidRequest {
95                error:   "Invalid message".into(),
96                request: msg.clone(),
97            }),
98        }
99    }
100}
pyth_sdk_cw/testing.rs

pyth_sdk_cw/
testing.rs
