kabu_node_debug_provider/
anvilprovider.rs1use alloy::{
2 primitives::{Address, Bytes, B256, U256, U64},
3 providers::{network::Ethereum, Network, Provider, RootProvider},
4 transports::TransportResult,
5};
6
7use crate::AnvilDebugProvider;
8
9pub(crate) fn convert_u64(r: U64) -> u64 {
10 r.to::<u64>()
11}
12
13pub trait AnvilProviderExt<N>
14where
15 N: Network,
16{
17 fn snapshot(&self) -> impl std::future::Future<Output = TransportResult<u64>> + Send;
18 fn revert(&self, snap_id: u64) -> impl std::future::Future<Output = TransportResult<bool>> + Send;
19
20 fn mine(&self) -> impl std::future::Future<Output = TransportResult<u64>> + Send;
21
22 fn set_automine(&self, to_mine: bool) -> impl std::future::Future<Output = TransportResult<()>> + Send;
23
24 fn set_code(&self, address: Address, code: Bytes) -> impl std::future::Future<Output = TransportResult<()>> + Send;
25
26 fn set_balance(&self, address: Address, balance: U256) -> impl std::future::Future<Output = TransportResult<()>> + Send;
27
28 fn set_storage(&self, address: Address, cell: B256, value: B256) -> impl std::future::Future<Output = TransportResult<bool>> + Send;
30}
31
32impl<PN, PA, N> AnvilProviderExt<N> for AnvilDebugProvider<PN, PA, N>
33where
34 N: Network,
35 PN: Provider<N> + Send + Sync + Clone + 'static,
36 PA: Provider<N> + Send + Sync + Clone + 'static,
37{
38 fn snapshot(&self) -> impl std::future::Future<Output = TransportResult<u64>> + Send {
39 self.anvil().client().request("evm_snapshot", ()).map_resp(convert_u64)
40 }
41 fn revert(&self, snap_id: u64) -> impl std::future::Future<Output = TransportResult<bool>> + Send {
42 self.anvil().client().request("evm_revert", (U64::from(snap_id),))
43 }
44
45 fn mine(&self) -> impl std::future::Future<Output = TransportResult<u64>> + Send {
46 self.anvil().client().request("evm_mine", ()).map_resp(convert_u64)
47 }
48
49 fn set_automine(&self, to_mine: bool) -> impl std::future::Future<Output = TransportResult<()>> + Send {
50 self.anvil().client().request("evm_setAutomine", (to_mine,))
51 }
52
53 fn set_code(&self, address: Address, code: Bytes) -> impl std::future::Future<Output = TransportResult<()>> + Send {
54 self.anvil().client().request("anvil_setCode", (address, code))
55 }
56
57 fn set_balance(&self, address: Address, balance: U256) -> impl std::future::Future<Output = TransportResult<()>> + Send {
58 self.anvil().client().request("anvil_setBalance", (address, balance))
59 }
60
61 fn set_storage(&self, address: Address, cell: B256, value: B256) -> impl std::future::Future<Output = TransportResult<bool>> + Send {
62 self.anvil().client().request("anvil_setStorageAt", (address, cell, value))
63 }
64 }
68
69impl AnvilProviderExt<Ethereum> for RootProvider<Ethereum> {
70 fn snapshot(&self) -> impl std::future::Future<Output = TransportResult<u64>> + Send {
71 self.client().request("evm_snapshot", ()).map_resp(convert_u64)
72 }
73 fn revert(&self, snap_id: u64) -> impl std::future::Future<Output = TransportResult<bool>> + Send {
74 self.client().request("evm_revert", (U64::from(snap_id),))
75 }
76
77 fn mine(&self) -> impl std::future::Future<Output = TransportResult<u64>> + Send {
78 self.client().request("evm_mine", ()).map_resp(convert_u64)
79 }
80
81 fn set_automine(&self, to_mine: bool) -> impl std::future::Future<Output = TransportResult<()>> + Send {
82 self.client().request("evm_setAutomine", (to_mine,))
83 }
84
85 fn set_code(&self, address: Address, code: Bytes) -> impl std::future::Future<Output = TransportResult<()>> + Send {
86 self.client().request("anvil_setCode", (address, code))
87 }
88
89 fn set_balance(&self, address: Address, balance: U256) -> impl std::future::Future<Output = TransportResult<()>> + Send {
90 self.client().request("anvil_setBalance", (address, balance))
91 }
92
93 fn set_storage(&self, address: Address, cell: B256, value: B256) -> impl std::future::Future<Output = TransportResult<bool>> + Send {
94 self.client().request("anvil_setStorageAt", (address, cell, value))
95 }
96 }
100
101#[cfg(test)]
102mod test {
103 use alloy::node_bindings::Anvil;
104 use alloy::primitives::{B256, U256};
105 use alloy::rpc::types::BlockNumberOrTag;
106 use alloy_provider::ProviderBuilder;
107 use alloy_rpc_client::ClientBuilder;
108 use env_logger::Env as EnvLog;
109 use eyre::Result;
110 use std::sync::Arc;
111
112 use super::*;
113
114 #[tokio::test]
115 async fn test_storage() -> Result<()> {
116 let _ = env_logger::try_init_from_env(EnvLog::default().default_filter_or("info"));
117
118 let anvil = Anvil::new().try_spawn()?;
119
120 let client_anvil = ClientBuilder::default().http(anvil.endpoint_url());
121
122 let provider = ProviderBuilder::new().disable_recommended_fillers().connect_client(client_anvil);
123
124 let address: Address = Address::repeat_byte(0x12);
125
126 if let Err(e) = provider.set_storage(address, B256::from(U256::from(1)), B256::from(U256::from(2))).await {
127 panic!("{}", e);
128 }
129
130 let value = provider.get_storage_at(address, U256::from(1)).await?;
131 if value != U256::from(2) {
132 panic!("Incorrect value {value}");
133 }
134
135 Ok(())
136 }
137
138 #[tokio::test]
139 async fn test_anvil_commands() -> Result<()> {
140 let _ = env_logger::try_init_from_env(EnvLog::default().default_filter_or("info"));
141 let anvil = Anvil::new().try_spawn()?;
142
143 dotenvy::from_filename(".env.test").ok();
144 let node_url = std::env::var("MAINNET_HTTP").unwrap().to_string();
145
146 let test_node_url = anvil.endpoint_url();
147
148 let node_url = url::Url::parse(node_url.as_str())?;
149
150 let client_anvil = ClientBuilder::default().http(test_node_url);
151 let provider_anvil = ProviderBuilder::new().disable_recommended_fillers().connect_client(client_anvil);
152
153 let client_node = ClientBuilder::default().http(node_url);
154 let provider_node = ProviderBuilder::new().disable_recommended_fillers().connect_client(client_node);
155
156 let provider = AnvilDebugProvider::new(provider_node, provider_anvil, BlockNumberOrTag::Number(10));
157
158 let client = Arc::new(provider);
159
160 let snap = client.snapshot().await?;
161 let _ = client.revert(snap).await?;
162 client.set_automine(false).await?;
163 let _ = client.mine().await;
164 client.set_automine(true).await?;
165
166 Ok(())
167 }
168}