Skip to main content

crucible_test_context/
mock_oracles.rs

1//! Mock oracle builders for testing DeFi protocols
2//!
3//! Provides builders for creating mock Pyth price feed accounts.
4//! Both marginfi and klend use identical Pyth Solana Receiver (PriceUpdateV2) format.
5
6use 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
17// ============================================================================
18// Pyth Solana Receiver Types (PriceUpdateV2)
19// ============================================================================
20
21/// PriceUpdateV2 discriminator: sha256("account:PriceUpdateV2")[0..8]
22pub const PYTH_DISCRIMINATOR: [u8; 8] = [34, 241, 35, 99, 157, 126, 244, 205];
23
24/// Default Pyth Solana Receiver program ID
25/// rec5EKMGg6MxZYaMdyBps2bnnCNHi6KCYuQedA7GsAuW
26pub 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/// Pyth verification level for price updates
32#[derive(Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
33#[repr(u8)]
34pub enum VerificationLevel {
35    Partial { num_signatures: u8 },
36    Full,
37}
38
39/// Pyth price feed message containing price data
40#[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/// Pyth PriceUpdateV2 account structure (Pyth Solana Receiver format)
54#[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
63// ============================================================================
64// Mock Pyth Oracle Builder
65// ============================================================================
66
67/// Builder for creating mock Pyth price feed accounts
68///
69/// # Example
70/// ```ignore
71/// let oracle = ctx.create_mock_pyth_oracle()
72///     .price(100_00000000)  // $100 with 8 decimals
73///     .exponent(-8)
74///     .confidence(100_000)
75///     .build()?;
76/// ```
77pub 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    /// Set the price value (in smallest units based on exponent)
101    /// For example, $100 with exponent -8 = 100_00000000
102    pub fn price(mut self, price: i64) -> Self {
103        self.price = price;
104        self
105    }
106
107    /// Set the price exponent (typically -8 for USD prices)
108    pub fn exponent(mut self, exp: i32) -> Self {
109        self.exponent = exp;
110        self
111    }
112
113    /// Set the confidence interval
114    pub fn confidence(mut self, conf: u64) -> Self {
115        self.confidence = conf;
116        self
117    }
118
119    /// Set the publish time (defaults to current system time)
120    pub fn publish_time(mut self, time: i64) -> Self {
121        self.publish_time = Some(time);
122        self
123    }
124
125    /// Set a custom feed ID (defaults to oracle pubkey bytes)
126    pub fn feed_id(mut self, id: [u8; 32]) -> Self {
127        self.feed_id = Some(id);
128        self
129    }
130
131    /// Override the Pyth program ID (defaults to Pyth Solana Receiver)
132    pub fn program_id(mut self, id: Pubkey) -> Self {
133        self.program_id = id;
134        self
135    }
136
137    /// Build and create the mock Pyth oracle account
138    /// Returns the oracle's public key
139    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        // Serialize: discriminator + borsh data
168        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
183// ============================================================================
184// TestContext Helper Methods for Pyth Oracles
185// ============================================================================
186
187impl TestContext {
188    /// Create a mock Pyth oracle builder
189    ///
190    /// # Example
191    /// ```ignore
192    /// let oracle = ctx.create_mock_pyth_oracle()
193    ///     .price(100_00000000)  // $100
194    ///     .exponent(-8)
195    ///     .build()?;
196    /// ```
197    pub fn create_mock_pyth_oracle(&mut self) -> MockPythOracleBuilder<'_> {
198        MockPythOracleBuilder::new(self)
199    }
200
201    /// Update the price on an existing Pyth oracle
202    ///
203    /// # Arguments
204    /// * `oracle` - The oracle account pubkey
205    /// * `price` - New price value
206    /// * `exponent` - Price exponent (typically -8)
207    pub fn update_pyth_price(&mut self, oracle: &Pubkey, price: i64, exponent: i32) -> Result<()> {
208        let account = self.read_account(oracle)?;
209
210        // Skip discriminator (8 bytes), deserialize, update, reserialize
211        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        // Reserialize
227        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    /// Refresh a Pyth oracle's timestamp and slot to make it "fresh"
240    /// Use this before operations that check oracle staleness
241    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}