kabu_node_debug_provider/
anvilprovider.rs

1use 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 get_storage(&self, address: Address, cell: U256) -> impl std::future::Future<Output=TransportResult<U256>> + Send;
29    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    /*    fn get_storage(&self, address: Address, cell: U256) -> impl std::future::Future<Output=TransportResult<U256>> + Send {
65        self.anvil().client().request("anvil_getStorageAt", (address, cell))
66    }*/
67}
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    /*    fn get_storage(&self, address: Address, cell: U256) -> impl std::future::Future<Output=TransportResult<U256>> + Send {
97        self.anvil().client().request("anvil_getStorageAt", (address, cell))
98    }*/
99}
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}