use std::collections::HashMap;
use async_trait::async_trait;
use thiserror::Error;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountData {
pub data: Vec<u8>,
pub owner: [u8; 32],
pub lamports: u64,
}
#[derive(Debug, Error)]
pub enum UpstreamError {
#[error("upstream transport error: {0}")]
Transport(String),
#[error("upstream returned an RPC error: {0}")]
Rpc(String),
#[error("no fixture stub registered for RPC method '{method}'")]
MethodNotStubbed { method: String },
#[error("upstream request timed out after {millis}ms")]
Timeout { millis: u64 },
}
pub type UpstreamResult<T> = Result<T, UpstreamError>;
#[async_trait]
pub trait UpstreamClient: Send + Sync {
async fn rpc_call(&self, method: &str, params: serde_json::Value) -> UpstreamResult<Vec<u8>>;
async fn get_account(&self, address: &str) -> UpstreamResult<Option<AccountData>>;
}
type FixtureRpcHandler =
Box<dyn Fn(&serde_json::Value) -> UpstreamResult<serde_json::Value> + Send + Sync>;
pub struct FixtureUpstream {
accounts: HashMap<String, AccountData>,
rpc_responses: HashMap<String, FixtureRpcHandler>,
}
impl FixtureUpstream {
#[must_use]
pub fn new() -> Self {
Self {
accounts: HashMap::new(),
rpc_responses: HashMap::new(),
}
}
#[must_use]
pub fn with_account(mut self, address: impl Into<String>, data: AccountData) -> Self {
self.accounts.insert(address.into(), data);
self
}
#[must_use]
pub fn with_method<F>(mut self, method: impl Into<String>, handler: F) -> Self
where
F: Fn(&serde_json::Value) -> UpstreamResult<serde_json::Value> + Send + Sync + 'static,
{
self.rpc_responses.insert(method.into(), Box::new(handler));
self
}
}
impl Default for FixtureUpstream {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl UpstreamClient for FixtureUpstream {
async fn rpc_call(&self, method: &str, params: serde_json::Value) -> UpstreamResult<Vec<u8>> {
if let Some(handler) = self.rpc_responses.get(method) {
let value = handler(¶ms)?;
return serde_json::to_vec(&value)
.map_err(|e| UpstreamError::Transport(format!("serialize fixture result: {e}")));
}
Err(UpstreamError::MethodNotStubbed {
method: method.to_string(),
})
}
async fn get_account(&self, address: &str) -> UpstreamResult<Option<AccountData>> {
Ok(self.accounts.get(address).cloned())
}
}