1use anyhow::Result;
7use borsh::{BorshDeserialize, BorshSerialize};
8use solana_account::Account;
9use solana_keypair::Keypair;
10use solana_pubkey::Pubkey;
11use solana_signer::Signer;
12use std::time::{SystemTime, UNIX_EPOCH};
13
14use crate::account_builders::AccountBuilderBase;
15use crate::TestContext;
16
17pub const PYTH_DISCRIMINATOR: [u8; 8] = [34, 241, 35, 99, 157, 126, 244, 205];
23
24pub const DEFAULT_PYTH_RECEIVER_ID: Pubkey = Pubkey::new_from_array([
27 0x02, 0xe1, 0xae, 0xce, 0x70, 0xcc, 0x1b, 0xac, 0x7a, 0x72, 0xa9, 0x36, 0x74, 0xe4, 0x5a, 0x7b,
28 0xe1, 0xa8, 0xbd, 0x5a, 0x03, 0xbd, 0x7c, 0x50, 0xfd, 0x3f, 0xa2, 0xc5, 0xa4, 0x92, 0x88, 0x28,
29]);
30
31#[derive(Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
33#[repr(u8)]
34pub enum VerificationLevel {
35 Partial { num_signatures: u8 },
36 Full,
37}
38
39#[derive(Clone, Copy, BorshSerialize, BorshDeserialize)]
41#[repr(C)]
42pub struct PriceFeedMessage {
43 pub feed_id: [u8; 32],
44 pub price: i64,
45 pub conf: u64,
46 pub exponent: i32,
47 pub publish_time: i64,
48 pub prev_publish_time: i64,
49 pub ema_price: i64,
50 pub ema_conf: u64,
51}
52
53#[derive(Clone, Copy, BorshSerialize, BorshDeserialize)]
55#[repr(C)]
56pub struct PriceUpdateV2 {
57 pub write_authority: Pubkey,
58 pub verification_level: VerificationLevel,
59 pub price_message: PriceFeedMessage,
60 pub posted_slot: u64,
61}
62
63pub struct MockPythOracleBuilder<'a> {
78 ctx: &'a mut TestContext,
79 price: i64,
80 exponent: i32,
81 confidence: u64,
82 publish_time: Option<i64>,
83 feed_id: Option<[u8; 32]>,
84 program_id: Pubkey,
85}
86
87impl<'a> MockPythOracleBuilder<'a> {
88 pub fn new(ctx: &'a mut TestContext) -> Self {
89 Self {
90 ctx,
91 price: 0,
92 exponent: -8,
93 confidence: 100_000,
94 publish_time: None,
95 feed_id: None,
96 program_id: DEFAULT_PYTH_RECEIVER_ID,
97 }
98 }
99
100 pub fn price(mut self, price: i64) -> Self {
103 self.price = price;
104 self
105 }
106
107 pub fn exponent(mut self, exp: i32) -> Self {
109 self.exponent = exp;
110 self
111 }
112
113 pub fn confidence(mut self, conf: u64) -> Self {
115 self.confidence = conf;
116 self
117 }
118
119 pub fn publish_time(mut self, time: i64) -> Self {
121 self.publish_time = Some(time);
122 self
123 }
124
125 pub fn feed_id(mut self, id: [u8; 32]) -> Self {
127 self.feed_id = Some(id);
128 self
129 }
130
131 pub fn program_id(mut self, id: Pubkey) -> Self {
133 self.program_id = id;
134 self
135 }
136
137 pub fn build(self) -> Result<Pubkey> {
140 let oracle_keypair = Keypair::new();
141 let oracle_pubkey = oracle_keypair.pubkey();
142
143 let current_time = SystemTime::now()
144 .duration_since(UNIX_EPOCH)
145 .unwrap()
146 .as_secs() as i64;
147
148 let publish_time = self.publish_time.unwrap_or(current_time);
149 let feed_id = self.feed_id.unwrap_or(oracle_pubkey.to_bytes());
150
151 let price_update = PriceUpdateV2 {
152 write_authority: Pubkey::default(),
153 verification_level: VerificationLevel::Full,
154 price_message: PriceFeedMessage {
155 feed_id,
156 price: self.price,
157 conf: self.confidence,
158 exponent: self.exponent,
159 publish_time,
160 prev_publish_time: publish_time - 1,
161 ema_price: self.price,
162 ema_conf: self.confidence,
163 },
164 posted_slot: self.ctx.slot(),
165 };
166
167 let mut data = PYTH_DISCRIMINATOR.to_vec();
169 price_update.serialize(&mut data)?;
170
171 self.ctx
172 .create_account()
173 .pubkey(oracle_pubkey)
174 .owner(self.program_id)
175 .lamports(1_000_000_000)
176 .data(&data)
177 .create()?;
178
179 Ok(oracle_pubkey)
180 }
181}
182
183impl TestContext {
188 pub fn create_mock_pyth_oracle(&mut self) -> MockPythOracleBuilder<'_> {
198 MockPythOracleBuilder::new(self)
199 }
200
201 pub fn update_pyth_price(&mut self, oracle: &Pubkey, price: i64, exponent: i32) -> Result<()> {
208 let account = self.read_account(oracle)?;
209
210 let mut price_update: PriceUpdateV2 =
212 BorshDeserialize::deserialize(&mut &account.data[8..])?;
213
214 let current_time = SystemTime::now()
215 .duration_since(UNIX_EPOCH)
216 .unwrap()
217 .as_secs() as i64;
218
219 price_update.price_message.price = price;
220 price_update.price_message.exponent = exponent;
221 price_update.price_message.ema_price = price;
222 price_update.price_message.prev_publish_time = price_update.price_message.publish_time;
223 price_update.price_message.publish_time = current_time;
224 price_update.posted_slot = self.slot();
225
226 let mut new_data = PYTH_DISCRIMINATOR.to_vec();
228 price_update.serialize(&mut new_data)?;
229
230 self.write_account(
231 oracle,
232 Account {
233 data: new_data,
234 ..account
235 },
236 )
237 }
238
239 pub fn refresh_pyth_oracle(&mut self, oracle: &Pubkey) -> Result<()> {
242 let account = self.read_account(oracle)?;
243
244 let mut price_update: PriceUpdateV2 =
245 BorshDeserialize::deserialize(&mut &account.data[8..])?;
246
247 let current_time = SystemTime::now()
248 .duration_since(UNIX_EPOCH)
249 .unwrap()
250 .as_secs() as i64;
251
252 price_update.price_message.prev_publish_time = price_update.price_message.publish_time;
253 price_update.price_message.publish_time = current_time;
254 price_update.posted_slot = self.slot();
255
256 let mut new_data = PYTH_DISCRIMINATOR.to_vec();
257 price_update.serialize(&mut new_data)?;
258
259 self.write_account(
260 oracle,
261 Account {
262 data: new_data,
263 ..account
264 },
265 )
266 }
267}