atlas_rpc_client/
mock_sender.rs

1//! A nonblocking [`RpcSender`] used for unit testing [`RpcClient`](crate::rpc_client::RpcClient).
2
3use {
4    crate::rpc_sender::*,
5    async_trait::async_trait,
6    base64::{prelude::BASE64_STANDARD, Engine},
7    serde_json::{json, Number, Value},
8    solana_account_decoder_client_types::{UiAccount, UiAccountData, UiAccountEncoding},
9    solana_clock::{Slot, UnixTimestamp},
10    solana_epoch_info::EpochInfo,
11    solana_epoch_schedule::EpochSchedule,
12    solana_instruction::{error::InstructionError, TRANSACTION_LEVEL_STACK_HEIGHT},
13    solana_message::MessageHeader,
14    solana_pubkey::Pubkey,
15    solana_rpc_client_api::{
16        client_error::Result,
17        config::RpcBlockProductionConfig,
18        request::RpcRequest,
19        response::{
20            Response, RpcAccountBalance, RpcBlockProduction, RpcBlockProductionRange, RpcBlockhash,
21            RpcConfirmedTransactionStatusWithSignature, RpcContactInfo, RpcIdentity,
22            RpcInflationGovernor, RpcInflationRate, RpcInflationReward, RpcKeyedAccount,
23            RpcPerfSample, RpcPrioritizationFee, RpcResponseContext, RpcSimulateTransactionResult,
24            RpcSnapshotSlotInfo, RpcSupply, RpcVersionInfo, RpcVoteAccountInfo,
25            RpcVoteAccountStatus,
26        },
27    },
28    solana_signature::Signature,
29    solana_transaction::{versioned::TransactionVersion, Transaction},
30    solana_transaction_error::{TransactionError, TransactionResult},
31    solana_transaction_status_client_types::{
32        option_serializer::OptionSerializer, EncodedConfirmedBlock,
33        EncodedConfirmedTransactionWithStatusMeta, EncodedTransaction,
34        EncodedTransactionWithStatusMeta, Rewards, TransactionBinaryEncoding,
35        TransactionConfirmationStatus, TransactionStatus, UiCompiledInstruction, UiMessage,
36        UiRawMessage, UiTransaction, UiTransactionStatusMeta,
37    },
38    solana_version::Version,
39    std::{
40        collections::{HashMap, VecDeque},
41        net::SocketAddr,
42        str::FromStr,
43        sync::RwLock,
44    },
45};
46
47pub const PUBKEY: &str = "7RoSF9fUmdphVCpabEoefH81WwrW7orsWonXWqTXkKV8";
48
49pub type Mocks = HashMap<RpcRequest, Value>;
50
51impl From<Mocks> for MocksMap {
52    fn from(mocks: Mocks) -> Self {
53        let mut map = HashMap::new();
54        for (key, value) in mocks {
55            map.insert(key, [value].into());
56        }
57        MocksMap(map)
58    }
59}
60
61#[derive(Default, Clone)]
62pub struct MocksMap(pub HashMap<RpcRequest, VecDeque<Value>>);
63
64impl FromIterator<(RpcRequest, Value)> for MocksMap {
65    fn from_iter<T: IntoIterator<Item = (RpcRequest, Value)>>(iter: T) -> Self {
66        let mut map = MocksMap::default();
67        for (request, value) in iter {
68            map.insert(request, value);
69        }
70        map
71    }
72}
73
74impl MocksMap {
75    pub fn insert(&mut self, request: RpcRequest, value: Value) {
76        let queue = self.0.entry(request).or_default();
77        queue.push_back(value)
78    }
79
80    pub fn pop_front_with_request(&mut self, request: &RpcRequest) -> Option<Value> {
81        self.0.get_mut(request).and_then(|queue| queue.pop_front())
82    }
83}
84
85pub struct MockSender {
86    mocks: RwLock<MocksMap>,
87    url: String,
88}
89
90/// An [`RpcSender`] used for unit testing [`RpcClient`](crate::rpc_client::RpcClient).
91///
92/// This is primarily for internal use.
93///
94/// Unless directed otherwise, it will generally return a reasonable default
95/// response, at least for [`RpcRequest`] values for which responses have been
96/// implemented.
97///
98/// The behavior can be customized in two ways:
99///
100/// 1) The `url` constructor argument is not actually a URL, but a simple string
101///    directive that changes `MockSender`s behavior in specific scenarios.
102///
103///    If `url` is "fails" then any call to `send` will return `Ok(Value::Null)`.
104///
105///    It is customary to set the `url` to "succeeds" for mocks that should
106///    return successfully, though this value is not actually interpreted.
107///
108///    Other possible values of `url` are specific to different `RpcRequest`
109///    values. Read the implementation for specifics.
110///
111/// 2) Custom responses can be configured by providing [`Mocks`] to the
112///    [`MockSender::new_with_mocks`] constructor. This type is a [`HashMap`]
113///    from [`RpcRequest`] to a JSON [`Value`] response, Any entries in this map
114///    override the default behavior for the given request.
115impl MockSender {
116    pub fn new<U: ToString>(url: U) -> Self {
117        Self::new_with_mocks(url, Mocks::default())
118    }
119
120    pub fn new_with_mocks<U: ToString>(url: U, mocks: Mocks) -> Self {
121        Self {
122            url: url.to_string(),
123            mocks: RwLock::new(MocksMap::from(mocks)),
124        }
125    }
126
127    pub fn new_with_mocks_map<U: ToString>(url: U, mocks: MocksMap) -> Self {
128        Self {
129            url: url.to_string(),
130            mocks: RwLock::new(mocks),
131        }
132    }
133}
134
135#[async_trait]
136impl RpcSender for MockSender {
137    fn get_transport_stats(&self) -> RpcTransportStats {
138        RpcTransportStats::default()
139    }
140
141    async fn send(
142        &self,
143        request: RpcRequest,
144        params: serde_json::Value,
145    ) -> Result<serde_json::Value> {
146        if let Some(value) = self.mocks.write().unwrap().pop_front_with_request(&request) {
147            return Ok(value);
148        }
149        if self.url == "fails" {
150            return Ok(Value::Null);
151        }
152
153        let method = &request.build_request_json(42, params.clone())["method"];
154
155        let val = match method.as_str().unwrap() {
156            "getAccountInfo" => serde_json::to_value(Response {
157                context: RpcResponseContext { slot: 1, api_version: None },
158                value: Value::Null,
159            })?,
160            "getBalance" => serde_json::to_value(Response {
161                context: RpcResponseContext { slot: 1, api_version: None },
162                value: Value::Number(Number::from(50)),
163            })?,
164            "getEpochInfo" => serde_json::to_value(EpochInfo {
165                epoch: 1,
166                slot_index: 2,
167                slots_in_epoch: 32,
168                absolute_slot: 34,
169                block_height: 34,
170                transaction_count: Some(123),
171            })?,
172            "getSignatureStatuses" => {
173                let status: TransactionResult<()> = if self.url == "account_in_use" {
174                    Err(TransactionError::AccountInUse)
175                } else if self.url == "instruction_error" {
176                    Err(TransactionError::InstructionError(
177                        0,
178                        InstructionError::UninitializedAccount,
179                    ))
180                } else {
181                    Ok(())
182                };
183                let status = if self.url == "sig_not_found" {
184                    None
185                } else {
186                    let err = status.clone().err();
187                    Some(TransactionStatus {
188                        status,
189                        slot: 1,
190                        confirmations: None,
191                        err,
192                        confirmation_status: Some(TransactionConfirmationStatus::Finalized),
193                    })
194                };
195                let statuses: Vec<Option<TransactionStatus>> = params.as_array().unwrap()[0]
196                    .as_array()
197                    .unwrap()
198                    .iter()
199                    .map(|_| status.clone())
200                    .collect();
201                serde_json::to_value(Response {
202                    context: RpcResponseContext { slot: 1, api_version: None },
203                    value: statuses,
204                })?
205            }
206            "getTransaction" => serde_json::to_value(EncodedConfirmedTransactionWithStatusMeta {
207                slot: 2,
208                transaction: EncodedTransactionWithStatusMeta {
209                    version: Some(TransactionVersion::LEGACY),
210                    transaction: EncodedTransaction::Json(
211                        UiTransaction {
212                            signatures: vec!["3AsdoALgZFuq2oUVWrDYhg2pNeaLJKPLf8hU2mQ6U8qJxeJ6hsrPVpMn9ma39DtfYCrDQSvngWRP8NnTpEhezJpE".to_string()],
213                            message: UiMessage::Raw(
214                                UiRawMessage {
215                                    header: MessageHeader {
216                                        num_required_signatures: 1,
217                                        num_readonly_signed_accounts: 0,
218                                        num_readonly_unsigned_accounts: 1,
219                                    },
220                                    account_keys: vec![
221                                        "C6eBmAXKg6JhJWkajGa5YRGUfG4YKXwbxF5Ufv7PtExZ".to_string(),
222                                        "2Gd5eoR5J4BV89uXbtunpbNhjmw3wa1NbRHxTHzDzZLX".to_string(),
223                                        "11111111111111111111111111111111".to_string(),
224                                    ],
225                                    recent_blockhash: "D37n3BSG71oUWcWjbZ37jZP7UfsxG2QMKeuALJ1PYvM6".to_string(),
226                                    instructions: vec![UiCompiledInstruction {
227                                        program_id_index: 2,
228                                        accounts: vec![0, 1],
229                                        data: "3Bxs49DitAvXtoDR".to_string(),
230                                        stack_height: Some(TRANSACTION_LEVEL_STACK_HEIGHT as u32),
231                                    }],
232                                    address_table_lookups: None,
233                                })
234                        }),
235                    meta: Some(UiTransactionStatusMeta {
236                            err: None,
237                            status: Ok(()),
238                            fee: 0,
239                            pre_balances: vec![499999999999999950, 50, 1],
240                            post_balances: vec![499999999999999950, 50, 1],
241                            inner_instructions: OptionSerializer::None,
242                            log_messages: OptionSerializer::None,
243                            pre_token_balances: OptionSerializer::None,
244                            post_token_balances: OptionSerializer::None,
245                            rewards: OptionSerializer::None,
246                            loaded_addresses: OptionSerializer::Skip,
247                            return_data: OptionSerializer::Skip,
248                            compute_units_consumed: OptionSerializer::Skip,
249                            cost_units: OptionSerializer::Skip,
250                        }),
251                },
252                block_time: Some(1628633791),
253            })?,
254            "getTransactionCount" => json![1234],
255            "getSlot" => json![0],
256            "getMaxShredInsertSlot" => json![0],
257            "requestAirdrop" => Value::String(Signature::from([8; 64]).to_string()),
258            "getHighestSnapshotSlot" => json!(RpcSnapshotSlotInfo {
259                full: 100,
260                incremental: Some(110),
261            }),
262            "getBlockHeight" => Value::Number(Number::from(1234)),
263            "getSlotLeaders" => json!([PUBKEY]),
264            "getBlockProduction" => {
265                if params.is_null() {
266                    json!(Response {
267                        context: RpcResponseContext { slot: 1, api_version: None },
268                        value: RpcBlockProduction {
269                            by_identity: HashMap::new(),
270                            range: RpcBlockProductionRange {
271                                first_slot: 1,
272                                last_slot: 2,
273                            },
274                        },
275                    })
276                } else {
277                    let config: Vec<RpcBlockProductionConfig> =
278                        serde_json::from_value(params).unwrap();
279                    let config = config[0].clone();
280                    let mut by_identity = HashMap::new();
281                    by_identity.insert(config.identity.unwrap(), (1, 123));
282                    let config_range = config.range.unwrap_or_default();
283
284                    json!(Response {
285                        context: RpcResponseContext { slot: 1, api_version: None },
286                        value: RpcBlockProduction {
287                            by_identity,
288                            range: RpcBlockProductionRange {
289                                first_slot: config_range.first_slot,
290                                last_slot: {
291                                    config_range.last_slot.unwrap_or(2)
292                                },
293                            },
294                        },
295                    })
296                }
297            }
298            "getStakeMinimumDelegation" => json!(Response {
299                context: RpcResponseContext { slot: 1, api_version: None },
300                value: 123_456_789,
301            }),
302            "getSupply" => json!(Response {
303                context: RpcResponseContext { slot: 1, api_version: None },
304                value: RpcSupply {
305                    total: 100000000,
306                    circulating: 50000,
307                    non_circulating: 20000,
308                    non_circulating_accounts: vec![PUBKEY.to_string()],
309                },
310            }),
311            "getLargestAccounts" => {
312                let rpc_account_balance = RpcAccountBalance {
313                    address: PUBKEY.to_string(),
314                    lamports: 10000,
315                };
316
317                json!(Response {
318                    context: RpcResponseContext { slot: 1, api_version: None },
319                    value: vec![rpc_account_balance],
320                })
321            }
322            "getVoteAccounts" => {
323                json!(RpcVoteAccountStatus {
324                    current: vec![],
325                    delinquent: vec![RpcVoteAccountInfo {
326                        vote_pubkey: PUBKEY.to_string(),
327                        node_pubkey: PUBKEY.to_string(),
328                        activated_stake: 0,
329                        commission: 0,
330                        epoch_vote_account: false,
331                        epoch_credits: vec![],
332                        last_vote: 0,
333                        root_slot: Slot::default(),
334                    }],
335                })
336            }
337            "sendTransaction" => {
338                let signature = if self.url == "malicious" {
339                    Signature::from([8; 64]).to_string()
340                } else {
341                    let tx_str = params.as_array().unwrap()[0].as_str().unwrap().to_string();
342                    let data = BASE64_STANDARD.decode(tx_str).unwrap();
343                    let tx: Transaction = bincode::deserialize(&data).unwrap();
344                    tx.signatures[0].to_string()
345                };
346                Value::String(signature)
347            }
348            "simulateTransaction" => serde_json::to_value(Response {
349                context: RpcResponseContext { slot: 1, api_version: None },
350                value: RpcSimulateTransactionResult {
351                    err: None,
352                    logs: None,
353                    accounts: None,
354                    units_consumed: None,
355                    loaded_accounts_data_size: None,
356                    return_data: None,
357                    inner_instructions: None,
358                    replacement_blockhash: None,
359                    fee: None,
360                    pre_balances: None,
361                    post_balances: None,
362                    pre_token_balances: None,
363                    post_token_balances: None,
364                    loaded_addresses: None,
365                }
366            })?,
367            "getMinimumBalanceForRentExemption" => json![20],
368            "getVersion" => {
369                let version = Version::default();
370                json!(RpcVersionInfo {
371                    solana_core: version.to_string(),
372                    feature_set: Some(version.feature_set),
373                })
374            }
375            "getLatestBlockhash" => serde_json::to_value(Response {
376                context: RpcResponseContext { slot: 1, api_version: None },
377                value: RpcBlockhash {
378                    blockhash: PUBKEY.to_string(),
379                    last_valid_block_height: 1234,
380                },
381            })?,
382            "getFeeForMessage" => serde_json::to_value(Response {
383                context: RpcResponseContext { slot: 1, api_version: None },
384                value: json!(Some(0)),
385            })?,
386            "getClusterNodes" => serde_json::to_value(vec![RpcContactInfo {
387                pubkey: PUBKEY.to_string(),
388                gossip: Some(SocketAddr::from(([10, 239, 6, 48], 8899))),
389                tvu: Some(SocketAddr::from(([10, 239, 6, 48], 8865))),
390                tpu: Some(SocketAddr::from(([10, 239, 6, 48], 8856))),
391                tpu_quic: Some(SocketAddr::from(([10, 239, 6, 48], 8862))),
392                tpu_forwards: Some(SocketAddr::from(([10, 239, 6, 48], 8857))),
393                tpu_forwards_quic: Some(SocketAddr::from(([10, 239, 6, 48], 8863))),
394                tpu_vote: Some(SocketAddr::from(([10, 239, 6, 48], 8870))),
395                serve_repair: Some(SocketAddr::from(([10, 239, 6, 48], 8880))),
396                rpc: Some(SocketAddr::from(([10, 239, 6, 48], 8899))),
397                pubsub: Some(SocketAddr::from(([10, 239, 6, 48], 8900))),
398                version: Some("1.0.0 c375ce1f".to_string()),
399                feature_set: None,
400                shred_version: None,
401            }])?,
402            "getBlock" => serde_json::to_value(EncodedConfirmedBlock {
403                previous_blockhash: "mfcyqEXB3DnHXki6KjjmZck6YjmZLvpAByy2fj4nh6B".to_string(),
404                blockhash: "3Eq21vXNB5s86c62bVuUfTeaMif1N2kUqRPBmGRJhyTA".to_string(),
405                parent_slot: 429,
406                transactions: vec![EncodedTransactionWithStatusMeta {
407                    transaction: EncodedTransaction::Binary(
408                        "ju9xZWuDBX4pRxX2oZkTjxU5jB4SSTgEGhX8bQ8PURNzyzqKMPPpNvWihx8zUe\
409                                 FfrbVNoAaEsNKZvGzAnTDy5bhNT9kt6KFCTBixpvrLCzg4M5UdFUQYrn1gdgjX\
410                                 pLHxcaShD81xBNaFDgnA2nkkdHnKtZt4hVSfKAmw3VRZbjrZ7L2fKZBx21CwsG\
411                                 hD6onjM2M3qZW5C8J6d1pj41MxKmZgPBSha3MyKkNLkAGFASK"
412                            .to_string(),
413                        TransactionBinaryEncoding::Base58,
414                    ),
415                    meta: None,
416                    version: Some(TransactionVersion::LEGACY),
417                }],
418                rewards: Rewards::new(),
419                num_partitions: None,
420                block_time: None,
421                block_height: Some(428),
422            })?,
423            "getBlocks" => serde_json::to_value(vec![1, 2, 3])?,
424            "getBlocksWithLimit" => serde_json::to_value(vec![1, 2, 3])?,
425            "getSignaturesForAddress" => {
426                serde_json::to_value(vec![RpcConfirmedTransactionStatusWithSignature {
427                    signature: crate::mock_sender_for_cli::SIGNATURE.to_string(),
428                    slot: 123,
429                    err: None,
430                    memo: None,
431                    block_time: None,
432                    confirmation_status: Some(TransactionConfirmationStatus::Finalized),
433                }])?
434            }
435            "getBlockTime" => serde_json::to_value(UnixTimestamp::default())?,
436            "getEpochSchedule" => serde_json::to_value(EpochSchedule::default())?,
437            "getRecentPerformanceSamples" => serde_json::to_value(vec![RpcPerfSample {
438                slot: 347873,
439                num_transactions: 125,
440                num_non_vote_transactions: Some(1),
441                num_slots: 123,
442                sample_period_secs: 60,
443            }])?,
444            "getRecentPrioritizationFees" => serde_json::to_value(vec![RpcPrioritizationFee {
445                slot: 123_456_789,
446                prioritization_fee: 10_000,
447            }])?,
448            "getIdentity" => serde_json::to_value(RpcIdentity {
449                identity: PUBKEY.to_string(),
450            })?,
451            "getInflationGovernor" => serde_json::to_value(
452                RpcInflationGovernor {
453                    initial: 0.08,
454                    terminal: 0.015,
455                    taper: 0.15,
456                    foundation: 0.05,
457                    foundation_term: 7.0,
458                })?,
459            "getInflationRate" => serde_json::to_value(
460                RpcInflationRate {
461                    total: 0.08,
462                    validator: 0.076,
463                    foundation: 0.004,
464                    epoch: 0,
465                })?,
466            "getInflationReward" => serde_json::to_value(vec![
467                Some(RpcInflationReward {
468                    epoch: 2,
469                    effective_slot: 224,
470                    amount: 2500,
471                    post_balance: 499999442500,
472                    commission: None,
473                })])?,
474            "minimumLedgerSlot" => json![123],
475            "getMaxRetransmitSlot" => json![123],
476            "getMultipleAccounts" => serde_json::to_value(Response {
477                context: RpcResponseContext { slot: 1, api_version: None },
478                value: vec![Value::Null, Value::Null]
479            })?,
480            "getProgramAccounts" => {
481                let pubkey = Pubkey::from_str(PUBKEY).unwrap();
482                serde_json::to_value(vec![
483                    RpcKeyedAccount {
484                        pubkey: PUBKEY.to_string(),
485                        account: mock_encoded_account(&pubkey)
486                    }
487                ])?
488            },
489            _ => Value::Null,
490        };
491        Ok(val)
492    }
493
494    fn url(&self) -> String {
495        format!("MockSender: {}", self.url)
496    }
497}
498
499pub(crate) fn mock_encoded_account(pubkey: &Pubkey) -> UiAccount {
500    UiAccount {
501        lamports: 1_000_000,
502        data: UiAccountData::Binary("".to_string(), UiAccountEncoding::Base64),
503        owner: pubkey.to_string(),
504        executable: false,
505        rent_epoch: 0,
506        space: Some(0),
507    }
508}
509
510#[cfg(test)]
511mod tests {
512    use {super::*, solana_account::Account, solana_account_decoder::encode_ui_account};
513
514    #[test]
515    fn test_mock_encoded_account() {
516        let pubkey = Pubkey::from_str(PUBKEY).unwrap();
517        let account = Account {
518            lamports: 1_000_000,
519            data: vec![],
520            owner: pubkey,
521            executable: false,
522            rent_epoch: 0,
523        };
524        let expected = encode_ui_account(&pubkey, &account, UiAccountEncoding::Base64, None, None);
525        assert_eq!(expected, mock_encoded_account(&pubkey));
526    }
527}