1use std::collections::HashMap;
2
3use alloy::primitives::U256;
4use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader};
5use tycho_common::{models::token::Token, Bytes};
6use tycho_ethereum::BytesCodec;
7
8use super::state::RocketpoolState;
9use crate::protocol::{
10 errors::InvalidSnapshotError,
11 models::{DecoderContext, TryFromWithBlock},
12};
13
14impl TryFromWithBlock<ComponentWithState, BlockHeader> for RocketpoolState {
15 type Error = InvalidSnapshotError;
16
17 async fn try_from_with_header(
20 snapshot: ComponentWithState,
21 _block: BlockHeader,
22 _account_balances: &HashMap<Bytes, HashMap<Bytes, Bytes>>,
23 _all_tokens: &HashMap<Bytes, Token>,
24 _decoder_context: &DecoderContext,
25 ) -> Result<Self, Self::Error> {
26 let total_eth = snapshot
27 .state
28 .attributes
29 .get("total_eth")
30 .map(U256::from_bytes)
31 .ok_or_else(|| InvalidSnapshotError::MissingAttribute("total_eth".to_string()))?;
32 let reth_supply = snapshot
33 .state
34 .attributes
35 .get("reth_supply")
36 .map(U256::from_bytes)
37 .ok_or_else(|| InvalidSnapshotError::MissingAttribute("reth_supply".to_string()))?;
38
39 let deposit_contract_balance = snapshot
40 .state
41 .attributes
42 .get("deposit_contract_balance")
43 .map(U256::from_bytes)
44 .ok_or_else(|| {
45 InvalidSnapshotError::MissingAttribute("deposit_contract_balance".to_string())
46 })?;
47
48 let reth_contract_liquidity = snapshot
49 .state
50 .attributes
51 .get("reth_contract_liquidity")
52 .map(U256::from_bytes)
53 .ok_or_else(|| {
54 InvalidSnapshotError::MissingAttribute("reth_contract_liquidity".to_string())
55 })?;
56
57 let deposits_enabled = snapshot
58 .state
59 .attributes
60 .get("deposits_enabled")
61 .map(|val| !U256::from_bytes(val).is_zero())
62 .ok_or_else(|| {
63 InvalidSnapshotError::MissingAttribute("deposits_enabled".to_string())
64 })?;
65
66 let deposit_assigning_enabled = snapshot
67 .state
68 .attributes
69 .get("deposit_assigning_enabled")
70 .map(|val| !U256::from_bytes(val).is_zero())
71 .ok_or_else(|| {
72 InvalidSnapshotError::MissingAttribute("deposit_assigning_enabled".to_string())
73 })?;
74
75 let deposit_fee = snapshot
76 .state
77 .attributes
78 .get("deposit_fee")
79 .map(U256::from_bytes)
80 .ok_or_else(|| InvalidSnapshotError::MissingAttribute("deposit_fee".to_string()))?;
81
82 let min_deposit_amount = snapshot
83 .state
84 .attributes
85 .get("min_deposit_amount")
86 .map(U256::from_bytes)
87 .ok_or_else(|| {
88 InvalidSnapshotError::MissingAttribute("min_deposit_amount".to_string())
89 })?;
90
91 let max_deposit_pool_size = snapshot
92 .state
93 .attributes
94 .get("max_deposit_pool_size")
95 .map(U256::from_bytes)
96 .ok_or_else(|| {
97 InvalidSnapshotError::MissingAttribute("max_deposit_pool_size".to_string())
98 })?;
99
100 let deposit_assign_maximum = snapshot
101 .state
102 .attributes
103 .get("deposit_assign_maximum")
104 .map(U256::from_bytes)
105 .ok_or_else(|| {
106 InvalidSnapshotError::MissingAttribute("deposit_assign_maximum".to_string())
107 })?;
108
109 let deposit_assign_socialised_maximum = snapshot
110 .state
111 .attributes
112 .get("deposit_assign_socialised_maximum")
113 .map(U256::from_bytes)
114 .ok_or_else(|| {
115 InvalidSnapshotError::MissingAttribute(
116 "deposit_assign_socialised_maximum".to_string(),
117 )
118 })?;
119
120 let queue_variable_start = snapshot
121 .state
122 .attributes
123 .get("queue_variable_start")
124 .map(U256::from_bytes)
125 .ok_or_else(|| {
126 InvalidSnapshotError::MissingAttribute("queue_variable_start".to_string())
127 })?;
128
129 let queue_variable_end = snapshot
130 .state
131 .attributes
132 .get("queue_variable_end")
133 .map(U256::from_bytes)
134 .ok_or_else(|| {
135 InvalidSnapshotError::MissingAttribute("queue_variable_end".to_string())
136 })?;
137
138 Ok(RocketpoolState::new(
139 reth_supply,
140 total_eth,
141 deposit_contract_balance,
142 reth_contract_liquidity,
143 deposit_fee,
144 deposits_enabled,
145 min_deposit_amount,
146 max_deposit_pool_size,
147 deposit_assigning_enabled,
148 deposit_assign_maximum,
149 deposit_assign_socialised_maximum,
150 queue_variable_start,
151 queue_variable_end,
152 ))
153 }
154}
155
156#[cfg(test)]
157mod tests {
158 use std::collections::HashMap;
159
160 use alloy::primitives::U256;
161 use rstest::rstest;
162 use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader};
163 use tycho_common::{dto::ResponseProtocolState, Bytes};
164
165 use super::super::state::RocketpoolState;
166 use crate::protocol::{
167 errors::InvalidSnapshotError,
168 models::{DecoderContext, TryFromWithBlock},
169 };
170
171 fn header() -> BlockHeader {
172 BlockHeader {
173 number: 1,
174 hash: Bytes::from(vec![0; 32]),
175 parent_hash: Bytes::from(vec![0; 32]),
176 revert: false,
177 timestamp: 1,
178 }
179 }
180
181 fn create_test_snapshot() -> ComponentWithState {
182 ComponentWithState {
183 state: ResponseProtocolState {
184 component_id: "Rocketpool".to_owned(),
185 attributes: HashMap::from([
186 (
187 "total_eth".to_string(),
188 Bytes::from(U256::from(100_000_000_000_000_000_000u128).to_be_bytes_vec()),
189 ),
190 (
191 "reth_supply".to_string(),
192 Bytes::from(U256::from(95_000_000_000_000_000_000u128).to_be_bytes_vec()),
193 ),
194 (
195 "deposit_contract_balance".to_string(),
196 Bytes::from(U256::from(50_000_000_000_000_000_000u128).to_be_bytes_vec()),
197 ), (
199 "reth_contract_liquidity".to_string(),
200 Bytes::from(U256::from(10_000_000_000_000_000_000u128).to_be_bytes_vec()),
201 ), ("deposits_enabled".to_string(), Bytes::from(vec![0x01])),
203 ("deposit_assigning_enabled".to_string(), Bytes::from(vec![0x01])),
204 (
205 "deposit_fee".to_string(),
206 Bytes::from(U256::from(5_000_000_000_000_000u128).to_be_bytes_vec()),
207 ), (
209 "min_deposit_amount".to_string(),
210 Bytes::from(U256::from(10_000_000_000_000_000u128).to_be_bytes_vec()),
211 ), (
213 "max_deposit_pool_size".to_string(),
214 Bytes::from(
215 U256::from(5_000_000_000_000_000_000_000u128).to_be_bytes_vec(),
216 ),
217 ), (
219 "deposit_assign_maximum".to_string(),
220 Bytes::from(U256::from(10u64).to_be_bytes_vec()),
221 ),
222 (
223 "deposit_assign_socialised_maximum".to_string(),
224 Bytes::from(U256::from(2u64).to_be_bytes_vec()),
225 ),
226 (
227 "queue_variable_start".to_string(),
228 Bytes::from(U256::from(100u64).to_be_bytes_vec()),
229 ),
230 (
231 "queue_variable_end".to_string(),
232 Bytes::from(U256::from(105u64).to_be_bytes_vec()),
233 ),
234 ]),
235 balances: HashMap::new(),
236 },
237 component: Default::default(),
238 component_tvl: None,
239 entrypoints: Vec::new(),
240 }
241 }
242
243 #[tokio::test]
244 async fn test_rocketpool_try_from() {
245 let snapshot = create_test_snapshot();
246
247 let result = RocketpoolState::try_from_with_header(
248 snapshot,
249 header(),
250 &HashMap::new(),
251 &HashMap::new(),
252 &DecoderContext::new(),
253 )
254 .await;
255
256 assert!(result.is_ok());
257 let state = result.unwrap();
258 assert_eq!(state.total_eth, U256::from(100_000_000_000_000_000_000u128));
259 assert_eq!(state.reth_supply, U256::from(95_000_000_000_000_000_000u128));
260 assert_eq!(state.deposit_contract_balance, U256::from(50_000_000_000_000_000_000u128));
261 assert_eq!(state.reth_contract_liquidity, U256::from(10_000_000_000_000_000_000u128));
262 assert!(state.deposits_enabled);
263 assert!(state.deposit_assigning_enabled);
264 assert_eq!(state.min_deposit_amount, U256::from(10_000_000_000_000_000u128));
265 assert_eq!(state.max_deposit_pool_size, U256::from(5_000_000_000_000_000_000_000u128));
266 assert_eq!(state.queue_variable_start, U256::from(100u64));
267 assert_eq!(state.queue_variable_end, U256::from(105u64));
268 }
269
270 #[tokio::test]
271 async fn test_rocketpool_try_from_deposits_disabled() {
272 let eth_address = Bytes::from(vec![0u8; 20]);
273
274 let snapshot = ComponentWithState {
275 state: ResponseProtocolState {
276 component_id: "Rocketpool".to_owned(),
277 attributes: HashMap::from([
278 ("total_eth".to_string(), Bytes::from(U256::from(100u64).to_be_bytes_vec())),
279 ("reth_supply".to_string(), Bytes::from(U256::from(100u64).to_be_bytes_vec())),
280 (
281 "deposit_contract_balance".to_string(),
282 Bytes::from(U256::from(50u64).to_be_bytes_vec()),
283 ),
284 (
285 "reth_contract_liquidity".to_string(),
286 Bytes::from(U256::from(10u64).to_be_bytes_vec()),
287 ),
288 ("deposits_enabled".to_string(), Bytes::from(vec![0x00])), ("deposit_assigning_enabled".to_string(), Bytes::from(vec![0x00])), ("deposit_fee".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
291 (
292 "min_deposit_amount".to_string(),
293 Bytes::from(U256::from(0u64).to_be_bytes_vec()),
294 ),
295 (
296 "max_deposit_pool_size".to_string(),
297 Bytes::from(U256::from(1000u64).to_be_bytes_vec()),
298 ),
299 (
300 "deposit_assign_maximum".to_string(),
301 Bytes::from(U256::from(0u64).to_be_bytes_vec()),
302 ),
303 (
304 "deposit_assign_socialised_maximum".to_string(),
305 Bytes::from(U256::from(0u64).to_be_bytes_vec()),
306 ),
307 (
308 "queue_full_start".to_string(),
309 Bytes::from(U256::from(0u64).to_be_bytes_vec()),
310 ),
311 ("queue_full_end".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
312 (
313 "queue_half_start".to_string(),
314 Bytes::from(U256::from(0u64).to_be_bytes_vec()),
315 ),
316 ("queue_half_end".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
317 (
318 "queue_variable_start".to_string(),
319 Bytes::from(U256::from(0u64).to_be_bytes_vec()),
320 ),
321 (
322 "queue_variable_end".to_string(),
323 Bytes::from(U256::from(0u64).to_be_bytes_vec()),
324 ),
325 ]),
326 balances: HashMap::from([(
327 eth_address,
328 Bytes::from(U256::from(50u64).to_be_bytes_vec()),
329 )]),
330 },
331 component: Default::default(),
332 component_tvl: None,
333 entrypoints: Vec::new(),
334 };
335
336 let result = RocketpoolState::try_from_with_header(
337 snapshot,
338 header(),
339 &HashMap::new(),
340 &HashMap::new(),
341 &DecoderContext::new(),
342 )
343 .await;
344
345 assert!(result.is_ok());
346 let state = result.unwrap();
347 assert!(!state.deposits_enabled);
348 assert!(!state.deposit_assigning_enabled);
349 }
350
351 #[tokio::test]
352 #[rstest]
353 #[case::missing_total_eth("total_eth")]
354 #[case::missing_reth_supply("reth_supply")]
355 #[case::missing_deposit_contract_balance("deposit_contract_balance")]
356 #[case::missing_reth_contract_liquidity("reth_contract_liquidity")]
357 #[case::missing_deposits_enabled("deposits_enabled")]
358 #[case::missing_deposit_assigning_enabled("deposit_assigning_enabled")]
359 #[case::missing_deposit_fee("deposit_fee")]
360 #[case::missing_min_deposit_amount("min_deposit_amount")]
361 #[case::missing_max_deposit_pool_size("max_deposit_pool_size")]
362 #[case::missing_deposit_assign_maximum("deposit_assign_maximum")]
363 #[case::missing_deposit_assign_socialised_maximum("deposit_assign_socialised_maximum")]
364 #[case::missing_queue_variable_start("queue_variable_start")]
365 #[case::missing_queue_variable_end("queue_variable_end")]
366 async fn test_rocketpool_try_from_missing_attribute(#[case] missing_attribute: &str) {
367 let eth_address = Bytes::from(vec![0u8; 20]);
368
369 let mut attributes = HashMap::from([
370 ("total_eth".to_string(), Bytes::from(U256::from(100u64).to_be_bytes_vec())),
371 ("reth_supply".to_string(), Bytes::from(U256::from(100u64).to_be_bytes_vec())),
372 (
373 "deposit_contract_balance".to_string(),
374 Bytes::from(U256::from(50u64).to_be_bytes_vec()),
375 ),
376 (
377 "reth_contract_liquidity".to_string(),
378 Bytes::from(U256::from(10u64).to_be_bytes_vec()),
379 ),
380 ("deposits_enabled".to_string(), Bytes::from(vec![0x01])),
381 ("deposit_assigning_enabled".to_string(), Bytes::from(vec![0x01])),
382 ("deposit_fee".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
383 ("min_deposit_amount".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
384 (
385 "max_deposit_pool_size".to_string(),
386 Bytes::from(U256::from(1000u64).to_be_bytes_vec()),
387 ),
388 ("deposit_assign_maximum".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
389 (
390 "deposit_assign_socialised_maximum".to_string(),
391 Bytes::from(U256::from(0u64).to_be_bytes_vec()),
392 ),
393 ("queue_full_start".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
394 ("queue_full_end".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
395 ("queue_half_start".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
396 ("queue_half_end".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
397 ("queue_variable_start".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
398 ("queue_variable_end".to_string(), Bytes::from(U256::from(0u64).to_be_bytes_vec())),
399 ]);
400 attributes.remove(missing_attribute);
401
402 let snapshot = ComponentWithState {
403 state: ResponseProtocolState {
404 component_id: "Rocketpool".to_owned(),
405 attributes,
406 balances: HashMap::from([(
407 eth_address,
408 Bytes::from(U256::from(50u64).to_be_bytes_vec()),
409 )]),
410 },
411 component: Default::default(),
412 component_tvl: None,
413 entrypoints: Vec::new(),
414 };
415
416 let result = RocketpoolState::try_from_with_header(
417 snapshot,
418 header(),
419 &HashMap::new(),
420 &HashMap::new(),
421 &DecoderContext::new(),
422 )
423 .await;
424
425 assert!(result.is_err());
426 assert!(matches!(
427 result.unwrap_err(),
428 InvalidSnapshotError::MissingAttribute(ref x) if x == missing_attribute
429 ));
430 }
431}