1use alloy::{
2 primitives::{Address, FixedBytes, U256},
3 rpc::types::Log,
4};
5use anyhow::{anyhow, Ok, Result};
6use std::any::Any;
7
8use crate::pool::base::TopicList;
9use crate::pool::base::{EventApplicable, PoolInterface, PoolType, PoolTypeTrait};
10
11#[derive(Debug, Clone)]
12pub struct MockPool {
13 address: Address,
14 token0: Address,
15 token1: Address,
16 reserve0: U256,
17 reserve1: U256,
18 fee: f64,
19 pool_type: PoolType,
20 last_updated: u64,
21}
22
23impl MockPool {
24 pub fn new(
25 address: Address,
26 token0: Address,
27 token1: Address,
28 reserve0: U256,
29 reserve1: U256,
30 fee: f64,
31 pool_type: PoolType,
32 ) -> Self {
33 let current_time = chrono::Utc::now().timestamp() as u64;
34 Self {
35 address,
36 token0,
37 token1,
38 reserve0,
39 reserve1,
40 fee,
41 pool_type,
42 last_updated: current_time,
43 }
44 }
45
46 pub fn new_boxed() -> Box<dyn PoolInterface + Send + Sync> {
48 Box::new(Self::new_v2(
49 Address::ZERO,
50 Address::ZERO,
51 Address::ZERO,
52 U256::from(1000000u64),
53 U256::from(1000000u64),
54 ))
55 }
56
57 pub fn new_v2(
58 address: Address,
59 token0: Address,
60 token1: Address,
61 reserve0: U256,
62 reserve1: U256,
63 ) -> Self {
64 Self::new(
65 address,
66 token0,
67 token1,
68 reserve0,
69 reserve1,
70 0.003,
71 PoolType::UniswapV2,
72 )
73 }
74
75 pub fn new_v3(
76 address: Address,
77 token0: Address,
78 token1: Address,
79 reserve0: U256,
80 reserve1: U256,
81 fee: f64,
82 ) -> Self {
83 Self::new(
84 address,
85 token0,
86 token1,
87 reserve0,
88 reserve1,
89 fee,
90 PoolType::UniswapV3,
91 )
92 }
93
94 pub fn set_reserves(&mut self, reserve0: U256, reserve1: U256) {
95 self.reserve0 = reserve0;
96 self.reserve1 = reserve1;
97 self.last_updated = chrono::Utc::now().timestamp() as u64;
98 }
99}
100
101impl PoolInterface for MockPool {
102 fn calculate_output(&self, token_in: &Address, amount_in: U256) -> Result<U256> {
103 if token_in == &self.token0 {
104 let amount_in_with_fee = amount_in * U256::from(997);
105 let numerator = amount_in_with_fee * self.reserve1;
106 let denominator = self.reserve0 * U256::from(1000) + amount_in_with_fee;
107 Ok(numerator / denominator)
108 } else if token_in == &self.token1 {
109 let amount_in_with_fee = amount_in * U256::from(997);
110 let numerator = amount_in_with_fee * self.reserve0;
111 let denominator = self.reserve1 * U256::from(1000) + amount_in_with_fee;
112 Ok(numerator / denominator)
113 } else {
114 Err(anyhow!("Token not in pool"))
115 }
116 }
117
118 fn calculate_input(&self, token_out: &Address, amount_out: U256) -> Result<U256> {
119 if token_out == &self.token0 {
120 let numerator = self.reserve1 * amount_out * U256::from(1000);
121 let denominator = (self.reserve0 - amount_out) * U256::from(997);
122 Ok((numerator / denominator) + U256::from(1))
123 } else if token_out == &self.token1 {
124 let numerator = self.reserve0 * amount_out * U256::from(1000);
125 let denominator = (self.reserve1 - amount_out) * U256::from(997);
126 Ok((numerator / denominator) + U256::from(1))
127 } else {
128 Err(anyhow!("Token not in pool"))
129 }
130 }
131
132 fn apply_swap(&mut self, token_in: &Address, amount_in: U256, amount_out: U256) -> Result<()> {
133 if token_in == &self.token0 {
134 if amount_out >= self.reserve1 {
135 return Err(anyhow!("Insufficient liquidity for swap"));
136 }
137 self.reserve0 += amount_in;
138 self.reserve1 -= amount_out;
139 } else if token_in == &self.token1 {
140 if amount_out >= self.reserve0 {
141 return Err(anyhow!("Insufficient liquidity for swap"));
142 }
143 self.reserve1 += amount_in;
144 self.reserve0 -= amount_out;
145 } else {
146 return Err(anyhow!("Token not in pool"));
147 }
148
149 self.last_updated = chrono::Utc::now().timestamp() as u64;
150 Ok(())
151 }
152
153 fn address(&self) -> Address {
154 self.address
155 }
156
157 fn tokens(&self) -> (Address, Address) {
158 (self.token0, self.token1)
159 }
160
161 fn log_summary(&self) -> String {
162 format!("Mock Pool")
163 }
164
165 fn fee(&self) -> f64 {
166 self.fee
167 }
168
169 fn fee_raw(&self) -> u64 {
170 self.fee as u64
171 }
172
173 fn id(&self) -> String {
174 format!(
175 "mock-{}-{}-{}-{:?}",
176 self.address, self.token0, self.token1, self.pool_type
177 )
178 }
179
180 fn contains_token(&self, token: &Address) -> bool {
181 *token == self.token0 || *token == self.token1
182 }
183
184 fn clone_box(&self) -> Box<dyn PoolInterface + Send + Sync> {
185 Box::new(self.clone())
186 }
187
188 fn as_any(&self) -> &dyn Any {
189 self
190 }
191
192 fn as_any_mut(&mut self) -> &mut dyn Any {
193 self
194 }
195}
196
197impl PoolTypeTrait for MockPool {
198 fn pool_type(&self) -> PoolType {
199 self.pool_type
200 }
201}
202
203impl EventApplicable for MockPool {
204 fn apply_log(&mut self, _event: &Log) -> Result<()> {
205 Ok(())
206 }
207}
208
209impl TopicList for MockPool {
210 fn topics() -> Vec<FixedBytes<32>> {
211 vec![]
212 }
213
214 fn profitable_topics() -> Vec<FixedBytes<32>> {
215 vec![]
216 }
217}
218
219#[cfg(test)]
220mod tests {
221 use super::*;
222
223 #[test]
224 fn test_mock_pool() {
225 let token0 = Address::ZERO;
226 let token1 = Address::repeat_byte(1);
227 let pool_addr = Address::repeat_byte(2);
228
229 let pool = MockPool::new_v2(
230 pool_addr,
231 token0,
232 token1,
233 U256::from(1000000),
234 U256::from(1000000),
235 );
236
237 assert_eq!(pool.address(), pool_addr);
238 assert_eq!(pool.tokens(), (token0, token1));
239 assert_eq!(pool.pool_type(), PoolType::UniswapV2);
240 assert_eq!(pool.fee(), 0.003);
241
242 let amount_out = pool.calculate_output(&token0, U256::from(1000)).unwrap();
243 assert!(amount_out > U256::ZERO);
244 }
245}