Skip to main content

tycho_common/
dto.rs

1//! Data Transfer Objects (or structs)
2//!
3//! These structs serve to serialise and deserialize messages between server and client, they should
4//! be very simple and ideally not contain any business logic.
5//!
6//! Structs in here implement utoipa traits so they can be used to derive an OpenAPI schema.
7#![allow(deprecated)]
8use std::{
9    collections::{BTreeMap, HashMap, HashSet},
10    fmt,
11    hash::{Hash, Hasher},
12};
13
14use chrono::{NaiveDateTime, Utc};
15use deepsize::{Context, DeepSizeOf};
16use serde::{de, Deserialize, Deserializer, Serialize};
17use strum_macros::{Display, EnumString};
18use thiserror::Error;
19use utoipa::{IntoParams, ToSchema};
20use uuid::Uuid;
21
22use crate::{
23    models::{
24        self, blockchain::BlockAggregatedChanges, Address, Balance, Code, ComponentId, StoreKey,
25        StoreVal,
26    },
27    serde_primitives::{
28        hex_bytes, hex_bytes_option, hex_hashmap_key, hex_hashmap_key_value, hex_hashmap_value,
29    },
30    Bytes,
31};
32
33/// Currently supported Blockchains
34#[derive(
35    Debug,
36    Clone,
37    Copy,
38    PartialEq,
39    Eq,
40    Hash,
41    Serialize,
42    Deserialize,
43    EnumString,
44    Display,
45    Default,
46    ToSchema,
47    DeepSizeOf,
48)]
49#[serde(rename_all = "lowercase")]
50#[strum(serialize_all = "lowercase")]
51pub enum Chain {
52    #[default]
53    Ethereum,
54    Starknet,
55    ZkSync,
56    Arbitrum,
57    Base,
58    Bsc,
59    Unichain,
60}
61
62impl From<models::contract::Account> for ResponseAccount {
63    fn from(value: models::contract::Account) -> Self {
64        ResponseAccount::new(
65            value.chain.into(),
66            value.address,
67            value.title,
68            value.slots,
69            value.native_balance,
70            value
71                .token_balances
72                .into_iter()
73                .map(|(k, v)| (k, v.balance))
74                .collect(),
75            value.code,
76            value.code_hash,
77            value.balance_modify_tx,
78            value.code_modify_tx,
79            value.creation_tx,
80        )
81    }
82}
83
84impl From<models::Chain> for Chain {
85    fn from(value: models::Chain) -> Self {
86        match value {
87            models::Chain::Ethereum => Chain::Ethereum,
88            models::Chain::Starknet => Chain::Starknet,
89            models::Chain::ZkSync => Chain::ZkSync,
90            models::Chain::Arbitrum => Chain::Arbitrum,
91            models::Chain::Base => Chain::Base,
92            models::Chain::Bsc => Chain::Bsc,
93            models::Chain::Unichain => Chain::Unichain,
94        }
95    }
96}
97
98#[derive(
99    Debug,
100    PartialEq,
101    Default,
102    Copy,
103    Clone,
104    Deserialize,
105    Serialize,
106    ToSchema,
107    EnumString,
108    Display,
109    DeepSizeOf,
110)]
111pub enum ChangeType {
112    #[default]
113    Update,
114    Deletion,
115    Creation,
116    Unspecified,
117}
118
119impl From<models::ChangeType> for ChangeType {
120    fn from(value: models::ChangeType) -> Self {
121        match value {
122            models::ChangeType::Update => ChangeType::Update,
123            models::ChangeType::Creation => ChangeType::Creation,
124            models::ChangeType::Deletion => ChangeType::Deletion,
125        }
126    }
127}
128
129impl ChangeType {
130    pub fn merge(&self, other: &Self) -> Self {
131        if matches!(self, Self::Creation) {
132            Self::Creation
133        } else {
134            *other
135        }
136    }
137}
138
139#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, Default)]
140pub struct ExtractorIdentity {
141    pub chain: Chain,
142    pub name: String,
143}
144
145impl ExtractorIdentity {
146    pub fn new(chain: Chain, name: &str) -> Self {
147        Self { chain, name: name.to_owned() }
148    }
149}
150
151impl fmt::Display for ExtractorIdentity {
152    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153        write!(f, "{}:{}", self.chain, self.name)
154    }
155}
156
157/// A command sent from the client to the server
158#[derive(Deserialize, Serialize, Debug, PartialEq, Eq)]
159#[serde(tag = "method", rename_all = "lowercase")]
160pub enum Command {
161    Subscribe {
162        extractor_id: ExtractorIdentity,
163        include_state: bool,
164        /// Enable zstd compression for messages in this subscription.
165        /// Defaults to false for backward compatibility.
166        #[serde(default)]
167        compression: bool,
168        /// Enables receiving partial block update messages in this subscription.
169        /// Defaults to false for backward compatibility.
170        #[serde(default)]
171        partial_blocks: bool,
172    },
173    Unsubscribe {
174        subscription_id: Uuid,
175    },
176}
177
178/// A easy serializable version of `models::error::WebsocketError`
179///
180/// This serves purely to transfer errors via websocket. It is meant to render
181/// similarly to the original struct but does not have server side debug information
182/// attached.
183///
184/// It should contain information needed to handle errors correctly on the client side.
185#[derive(Error, Debug, Serialize, Deserialize, Eq, PartialEq, Clone)]
186pub enum WebsocketError {
187    #[error("Extractor not found: {0}")]
188    ExtractorNotFound(ExtractorIdentity),
189
190    #[error("Subscription not found: {0}")]
191    SubscriptionNotFound(Uuid),
192
193    #[error("Failed to parse JSON: {1}, msg: {0}")]
194    ParseError(String, String),
195
196    #[error("Failed to subscribe to extractor: {0}")]
197    SubscribeError(ExtractorIdentity),
198
199    #[error("Failed to compress message for subscription: {0}, error: {1}")]
200    CompressionError(Uuid, String),
201}
202
203impl From<crate::models::error::WebsocketError> for WebsocketError {
204    fn from(value: crate::models::error::WebsocketError) -> Self {
205        match value {
206            crate::models::error::WebsocketError::ExtractorNotFound(eid) => {
207                Self::ExtractorNotFound(eid.into())
208            }
209            crate::models::error::WebsocketError::SubscriptionNotFound(sid) => {
210                Self::SubscriptionNotFound(sid)
211            }
212            crate::models::error::WebsocketError::ParseError(raw, error) => {
213                Self::ParseError(error.to_string(), raw)
214            }
215            crate::models::error::WebsocketError::SubscribeError(eid) => {
216                Self::SubscribeError(eid.into())
217            }
218            crate::models::error::WebsocketError::CompressionError(sid, error) => {
219                Self::CompressionError(sid, error.to_string())
220            }
221        }
222    }
223}
224
225/// A response sent from the server to the client
226#[derive(Deserialize, Serialize, Debug, PartialEq, Eq, Clone)]
227#[serde(tag = "method", rename_all = "lowercase")]
228pub enum Response {
229    NewSubscription { extractor_id: ExtractorIdentity, subscription_id: Uuid },
230    SubscriptionEnded { subscription_id: Uuid },
231    Error(WebsocketError),
232}
233
234/// A message sent from the server to the client
235#[allow(clippy::large_enum_variant)]
236#[derive(Serialize, Deserialize, Debug, Display, Clone)]
237#[serde(untagged)]
238pub enum WebSocketMessage {
239    BlockChanges { subscription_id: Uuid, deltas: BlockChanges },
240    Response(Response),
241}
242
243#[derive(Debug, PartialEq, Clone, Deserialize, Serialize, Default, ToSchema)]
244pub struct Block {
245    pub number: u64,
246    #[serde(with = "hex_bytes")]
247    pub hash: Bytes,
248    #[serde(with = "hex_bytes")]
249    pub parent_hash: Bytes,
250    pub chain: Chain,
251    pub ts: NaiveDateTime,
252}
253
254#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash, DeepSizeOf)]
255#[serde(deny_unknown_fields)]
256pub struct BlockParam {
257    #[schema(value_type=Option<String>)]
258    #[serde(with = "hex_bytes_option", default)]
259    pub hash: Option<Bytes>,
260    #[deprecated(
261        note = "The `chain` field is deprecated and will be removed in a future version."
262    )]
263    #[serde(default)]
264    pub chain: Option<Chain>,
265    #[serde(default)]
266    pub number: Option<i64>,
267}
268
269impl From<&Block> for BlockParam {
270    fn from(value: &Block) -> Self {
271        // The hash should uniquely identify a block across chains
272        BlockParam { hash: Some(value.hash.clone()), chain: None, number: None }
273    }
274}
275
276#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
277pub struct TokenBalances(#[serde(with = "hex_hashmap_key")] pub HashMap<Bytes, ComponentBalance>);
278
279impl From<HashMap<Bytes, ComponentBalance>> for TokenBalances {
280    fn from(value: HashMap<Bytes, ComponentBalance>) -> Self {
281        TokenBalances(value)
282    }
283}
284
285#[derive(Debug, PartialEq, Clone, Default, Deserialize, Serialize)]
286pub struct Transaction {
287    #[serde(with = "hex_bytes")]
288    pub hash: Bytes,
289    #[serde(with = "hex_bytes")]
290    pub block_hash: Bytes,
291    #[serde(with = "hex_bytes")]
292    pub from: Bytes,
293    #[serde(with = "hex_bytes_option")]
294    pub to: Option<Bytes>,
295    pub index: u64,
296}
297
298impl Transaction {
299    pub fn new(hash: Bytes, block_hash: Bytes, from: Bytes, to: Option<Bytes>, index: u64) -> Self {
300        Self { hash, block_hash, from, to, index }
301    }
302}
303
304/// A container for updates grouped by account/component.
305#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
306pub struct BlockChanges {
307    pub extractor: String,
308    pub chain: Chain,
309    pub block: Block,
310    pub finalized_block_height: u64,
311    pub revert: bool,
312    #[serde(with = "hex_hashmap_key", default)]
313    pub new_tokens: HashMap<Bytes, ResponseToken>,
314    #[serde(alias = "account_deltas", with = "hex_hashmap_key")]
315    pub account_updates: HashMap<Bytes, AccountUpdate>,
316    #[serde(alias = "state_deltas")]
317    pub state_updates: HashMap<String, ProtocolStateDelta>,
318    pub new_protocol_components: HashMap<String, ProtocolComponent>,
319    pub deleted_protocol_components: HashMap<String, ProtocolComponent>,
320    pub component_balances: HashMap<String, TokenBalances>,
321    pub account_balances: HashMap<Bytes, HashMap<Bytes, AccountBalance>>,
322    pub component_tvl: HashMap<String, f64>,
323    pub dci_update: DCIUpdate,
324    /// The index of the partial block. None if it's a full block.
325    #[serde(default, skip_serializing_if = "Option::is_none")]
326    pub partial_block_index: Option<u32>,
327}
328
329impl BlockChanges {
330    #[allow(clippy::too_many_arguments)]
331    pub fn new(
332        extractor: &str,
333        chain: Chain,
334        block: Block,
335        finalized_block_height: u64,
336        revert: bool,
337        account_updates: HashMap<Bytes, AccountUpdate>,
338        state_updates: HashMap<String, ProtocolStateDelta>,
339        new_protocol_components: HashMap<String, ProtocolComponent>,
340        deleted_protocol_components: HashMap<String, ProtocolComponent>,
341        component_balances: HashMap<String, HashMap<Bytes, ComponentBalance>>,
342        account_balances: HashMap<Bytes, HashMap<Bytes, AccountBalance>>,
343        dci_update: DCIUpdate,
344    ) -> Self {
345        BlockChanges {
346            extractor: extractor.to_owned(),
347            chain,
348            block,
349            finalized_block_height,
350            revert,
351            new_tokens: HashMap::new(),
352            account_updates,
353            state_updates,
354            new_protocol_components,
355            deleted_protocol_components,
356            component_balances: component_balances
357                .into_iter()
358                .map(|(k, v)| (k, v.into()))
359                .collect(),
360            account_balances,
361            component_tvl: HashMap::new(),
362            dci_update,
363            partial_block_index: None,
364        }
365    }
366
367    pub fn merge(mut self, other: Self) -> Self {
368        other
369            .account_updates
370            .into_iter()
371            .for_each(|(k, v)| {
372                self.account_updates
373                    .entry(k)
374                    .and_modify(|e| {
375                        e.merge(&v);
376                    })
377                    .or_insert(v);
378            });
379
380        other
381            .state_updates
382            .into_iter()
383            .for_each(|(k, v)| {
384                self.state_updates
385                    .entry(k)
386                    .and_modify(|e| {
387                        e.merge(&v);
388                    })
389                    .or_insert(v);
390            });
391
392        other
393            .component_balances
394            .into_iter()
395            .for_each(|(k, v)| {
396                self.component_balances
397                    .entry(k)
398                    .and_modify(|e| e.0.extend(v.0.clone()))
399                    .or_insert_with(|| v);
400            });
401
402        other
403            .account_balances
404            .into_iter()
405            .for_each(|(k, v)| {
406                self.account_balances
407                    .entry(k)
408                    .and_modify(|e| e.extend(v.clone()))
409                    .or_insert(v);
410            });
411
412        self.component_tvl
413            .extend(other.component_tvl);
414        self.new_protocol_components
415            .extend(other.new_protocol_components);
416        self.deleted_protocol_components
417            .extend(other.deleted_protocol_components);
418        self.revert = other.revert;
419        self.block = other.block;
420
421        self
422    }
423
424    pub fn get_block(&self) -> &Block {
425        &self.block
426    }
427
428    pub fn is_revert(&self) -> bool {
429        self.revert
430    }
431
432    pub fn filter_by_component<F: Fn(&str) -> bool>(&mut self, keep: F) {
433        self.state_updates
434            .retain(|k, _| keep(k));
435        self.component_balances
436            .retain(|k, _| keep(k));
437        self.component_tvl
438            .retain(|k, _| keep(k));
439    }
440
441    pub fn filter_by_contract<F: Fn(&Bytes) -> bool>(&mut self, keep: F) {
442        self.account_updates
443            .retain(|k, _| keep(k));
444        self.account_balances
445            .retain(|k, _| keep(k));
446    }
447
448    pub fn n_changes(&self) -> usize {
449        self.account_updates.len() + self.state_updates.len()
450    }
451
452    pub fn drop_state(&self) -> Self {
453        Self {
454            extractor: self.extractor.clone(),
455            chain: self.chain,
456            block: self.block.clone(),
457            finalized_block_height: self.finalized_block_height,
458            revert: self.revert,
459            new_tokens: self.new_tokens.clone(),
460            account_updates: HashMap::new(),
461            state_updates: HashMap::new(),
462            new_protocol_components: self.new_protocol_components.clone(),
463            deleted_protocol_components: self.deleted_protocol_components.clone(),
464            component_balances: self.component_balances.clone(),
465            account_balances: self.account_balances.clone(),
466            component_tvl: self.component_tvl.clone(),
467            dci_update: self.dci_update.clone(),
468            partial_block_index: self.partial_block_index,
469        }
470    }
471
472    pub fn is_partial_block(&self) -> bool {
473        self.partial_block_index.is_some()
474    }
475}
476
477impl From<models::blockchain::Block> for Block {
478    fn from(value: models::blockchain::Block) -> Self {
479        Self {
480            number: value.number,
481            hash: value.hash,
482            parent_hash: value.parent_hash,
483            chain: value.chain.into(),
484            ts: value.ts,
485        }
486    }
487}
488
489impl From<models::protocol::ComponentBalance> for ComponentBalance {
490    fn from(value: models::protocol::ComponentBalance) -> Self {
491        Self {
492            token: value.token,
493            balance: value.balance,
494            balance_float: value.balance_float,
495            modify_tx: value.modify_tx,
496            component_id: value.component_id,
497        }
498    }
499}
500
501impl From<models::contract::AccountBalance> for AccountBalance {
502    fn from(value: models::contract::AccountBalance) -> Self {
503        Self {
504            account: value.account,
505            token: value.token,
506            balance: value.balance,
507            modify_tx: value.modify_tx,
508        }
509    }
510}
511
512impl From<BlockAggregatedChanges> for BlockChanges {
513    fn from(value: BlockAggregatedChanges) -> Self {
514        Self {
515            extractor: value.extractor,
516            chain: value.chain.into(),
517            block: value.block.into(),
518            finalized_block_height: value.finalized_block_height,
519            revert: value.revert,
520            account_updates: value
521                .account_deltas
522                .into_iter()
523                .map(|(k, v)| (k, v.into()))
524                .collect(),
525            state_updates: value
526                .state_deltas
527                .into_iter()
528                .map(|(k, v)| (k, v.into()))
529                .collect(),
530            new_protocol_components: value
531                .new_protocol_components
532                .into_iter()
533                .map(|(k, v)| (k, v.into()))
534                .collect(),
535            deleted_protocol_components: value
536                .deleted_protocol_components
537                .into_iter()
538                .map(|(k, v)| (k, v.into()))
539                .collect(),
540            component_balances: value
541                .component_balances
542                .into_iter()
543                .map(|(component_id, v)| {
544                    let balances: HashMap<Bytes, ComponentBalance> = v
545                        .into_iter()
546                        .map(|(k, v)| (k, ComponentBalance::from(v)))
547                        .collect();
548                    (component_id, balances.into())
549                })
550                .collect(),
551            account_balances: value
552                .account_balances
553                .into_iter()
554                .map(|(k, v)| {
555                    (
556                        k,
557                        v.into_iter()
558                            .map(|(k, v)| (k, v.into()))
559                            .collect(),
560                    )
561                })
562                .collect(),
563            dci_update: value.dci_update.into(),
564            new_tokens: value
565                .new_tokens
566                .into_iter()
567                .map(|(k, v)| (k, v.into()))
568                .collect(),
569            component_tvl: value.component_tvl,
570            partial_block_index: value.partial_block_index,
571        }
572    }
573}
574
575#[derive(PartialEq, Serialize, Deserialize, Clone, Debug, ToSchema)]
576pub struct AccountUpdate {
577    #[serde(with = "hex_bytes")]
578    #[schema(value_type=Vec<String>)]
579    pub address: Bytes,
580    pub chain: Chain,
581    #[serde(with = "hex_hashmap_key_value")]
582    #[schema(value_type=HashMap<String, String>)]
583    pub slots: HashMap<Bytes, Bytes>,
584    #[serde(with = "hex_bytes_option")]
585    #[schema(value_type=Option<String>)]
586    pub balance: Option<Bytes>,
587    #[serde(with = "hex_bytes_option")]
588    #[schema(value_type=Option<String>)]
589    pub code: Option<Bytes>,
590    pub change: ChangeType,
591}
592
593impl AccountUpdate {
594    pub fn new(
595        address: Bytes,
596        chain: Chain,
597        slots: HashMap<Bytes, Bytes>,
598        balance: Option<Bytes>,
599        code: Option<Bytes>,
600        change: ChangeType,
601    ) -> Self {
602        Self { address, chain, slots, balance, code, change }
603    }
604
605    pub fn merge(&mut self, other: &Self) {
606        self.slots.extend(
607            other
608                .slots
609                .iter()
610                .map(|(k, v)| (k.clone(), v.clone())),
611        );
612        self.balance.clone_from(&other.balance);
613        self.code.clone_from(&other.code);
614        self.change = self.change.merge(&other.change);
615    }
616}
617
618impl From<models::contract::AccountDelta> for AccountUpdate {
619    fn from(value: models::contract::AccountDelta) -> Self {
620        let code = value.code().clone();
621        let change_type = value.change_type().into();
622        AccountUpdate::new(
623            value.address,
624            value.chain.into(),
625            value
626                .slots
627                .into_iter()
628                .map(|(k, v)| (k, v.unwrap_or_default()))
629                .collect(),
630            value.balance,
631            code,
632            change_type,
633        )
634    }
635}
636
637/// Represents the static parts of a protocol component.
638#[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize, ToSchema)]
639pub struct ProtocolComponent {
640    /// Unique identifier for this component
641    pub id: String,
642    /// Protocol system this component is part of
643    pub protocol_system: String,
644    /// Type of the protocol system
645    pub protocol_type_name: String,
646    pub chain: Chain,
647    /// Token addresses the component operates on
648    #[schema(value_type=Vec<String>)]
649    pub tokens: Vec<Bytes>,
650    /// Contract addresses involved in the components operations (may be empty for
651    /// native implementations)
652    #[serde(alias = "contract_addresses")]
653    #[schema(value_type=Vec<String>)]
654    pub contract_ids: Vec<Bytes>,
655    /// Constant attributes of the component
656    #[serde(with = "hex_hashmap_value")]
657    #[schema(value_type=HashMap<String, String>)]
658    pub static_attributes: HashMap<String, Bytes>,
659    /// Indicates if last change was update, create or delete (for internal use only).
660    #[serde(default)]
661    pub change: ChangeType,
662    /// Transaction hash which created this component
663    #[serde(with = "hex_bytes")]
664    #[schema(value_type=String)]
665    pub creation_tx: Bytes,
666    /// Date time of creation in UTC time
667    pub created_at: NaiveDateTime,
668}
669
670// Manual impl as `NaiveDateTime` structure referenced in `created_at` does not implement DeepSizeOf
671impl DeepSizeOf for ProtocolComponent {
672    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
673        self.id.deep_size_of_children(ctx) +
674            self.protocol_system
675                .deep_size_of_children(ctx) +
676            self.protocol_type_name
677                .deep_size_of_children(ctx) +
678            self.chain.deep_size_of_children(ctx) +
679            self.tokens.deep_size_of_children(ctx) +
680            self.contract_ids
681                .deep_size_of_children(ctx) +
682            self.static_attributes
683                .deep_size_of_children(ctx) +
684            self.change.deep_size_of_children(ctx) +
685            self.creation_tx
686                .deep_size_of_children(ctx)
687    }
688}
689
690impl<T> From<models::protocol::ProtocolComponent<T>> for ProtocolComponent
691where
692    T: Into<Address> + Clone,
693{
694    fn from(value: models::protocol::ProtocolComponent<T>) -> Self {
695        Self {
696            id: value.id,
697            protocol_system: value.protocol_system,
698            protocol_type_name: value.protocol_type_name,
699            chain: value.chain.into(),
700            tokens: value
701                .tokens
702                .into_iter()
703                .map(|t| t.into())
704                .collect(),
705            contract_ids: value.contract_addresses,
706            static_attributes: value.static_attributes,
707            change: value.change.into(),
708            creation_tx: value.creation_tx,
709            created_at: value.created_at,
710        }
711    }
712}
713
714#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
715pub struct ComponentBalance {
716    #[serde(with = "hex_bytes")]
717    pub token: Bytes,
718    pub balance: Bytes,
719    pub balance_float: f64,
720    #[serde(with = "hex_bytes")]
721    pub modify_tx: Bytes,
722    pub component_id: String,
723}
724
725#[derive(Debug, PartialEq, Clone, Default, Serialize, Deserialize, ToSchema)]
726/// Represents a change in protocol state.
727pub struct ProtocolStateDelta {
728    pub component_id: String,
729    #[schema(value_type=HashMap<String, String>)]
730    pub updated_attributes: HashMap<String, Bytes>,
731    pub deleted_attributes: HashSet<String>,
732}
733
734impl From<models::protocol::ProtocolComponentStateDelta> for ProtocolStateDelta {
735    fn from(value: models::protocol::ProtocolComponentStateDelta) -> Self {
736        Self {
737            component_id: value.component_id,
738            updated_attributes: value.updated_attributes,
739            deleted_attributes: value.deleted_attributes,
740        }
741    }
742}
743
744impl ProtocolStateDelta {
745    /// Merges 'other' into 'self'.
746    ///
747    ///
748    /// During merge of these deltas a special situation can arise when an attribute is present in
749    /// `self.deleted_attributes` and `other.update_attributes``. If we would just merge the sets
750    /// of deleted attributes or vice versa, it would be ambiguous and potential lead to a
751    /// deletion of an attribute that should actually be present, or retention of an actually
752    /// deleted attribute.
753    ///
754    /// This situation is handled the following way:
755    ///
756    ///     - If an attribute is deleted and in the next message recreated, it is removed from the
757    ///       set of deleted attributes and kept in updated_attributes. This way it's temporary
758    ///       deletion is never communicated to the final receiver.
759    ///     - If an attribute was updated and is deleted in the next message, it is removed from
760    ///       updated attributes and kept in deleted. This way the attributes temporary update (or
761    ///       potentially short-lived existence) before its deletion is never communicated to the
762    ///       final receiver.
763    pub fn merge(&mut self, other: &Self) {
764        // either updated and then deleted -> keep in deleted, remove from updated
765        self.updated_attributes
766            .retain(|k, _| !other.deleted_attributes.contains(k));
767
768        // or deleted and then updated/recreated -> remove from deleted and keep in updated
769        self.deleted_attributes.retain(|attr| {
770            !other
771                .updated_attributes
772                .contains_key(attr)
773        });
774
775        // simply merge updates
776        self.updated_attributes.extend(
777            other
778                .updated_attributes
779                .iter()
780                .map(|(k, v)| (k.clone(), v.clone())),
781        );
782
783        // simply merge deletions
784        self.deleted_attributes
785            .extend(other.deleted_attributes.iter().cloned());
786    }
787}
788
789/// Pagination parameter
790#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash, DeepSizeOf)]
791#[serde(deny_unknown_fields)]
792pub struct PaginationParams {
793    /// What page to retrieve
794    #[serde(default)]
795    pub page: i64,
796    /// How many results to return per page
797    #[serde(default)]
798    #[schema(default = 100)]
799    pub page_size: i64,
800}
801
802impl PaginationParams {
803    pub fn new(page: i64, page_size: i64) -> Self {
804        Self { page, page_size }
805    }
806}
807
808impl Default for PaginationParams {
809    fn default() -> Self {
810        PaginationParams { page: 0, page_size: 100 }
811    }
812}
813
814/// Defines pagination size limits for request types.
815///
816/// Different limits apply based on whether compression is enabled,
817/// as compressed responses can safely transfer more data.
818pub trait PaginationLimits {
819    /// Maximum page size when compression is enabled (e.g., zstd)
820    const MAX_PAGE_SIZE_COMPRESSED: i64;
821
822    /// Maximum page size when compression is disabled
823    const MAX_PAGE_SIZE_UNCOMPRESSED: i64;
824
825    /// Returns the effective maximum page size based on compression setting
826    fn effective_max_page_size(compression: bool) -> i64 {
827        if compression {
828            Self::MAX_PAGE_SIZE_COMPRESSED
829        } else {
830            Self::MAX_PAGE_SIZE_UNCOMPRESSED
831        }
832    }
833
834    /// Returns a reference to the pagination parameters
835    fn pagination(&self) -> &PaginationParams;
836}
837
838/// Macro to implement PaginationLimits for request types
839///
840/// When INCREASING these limits, ensure to immediately redeploy the servers.
841///
842/// Why: pagination limits are shared. Clients use these constants to set their max page size.
843/// When clients upgrade before servers, they request more than old servers allow and get errors.
844macro_rules! impl_pagination_limits {
845    ($type:ty, compressed = $comp:expr, uncompressed = $uncomp:expr) => {
846        impl $crate::dto::PaginationLimits for $type {
847            const MAX_PAGE_SIZE_COMPRESSED: i64 = $comp;
848            const MAX_PAGE_SIZE_UNCOMPRESSED: i64 = $uncomp;
849
850            fn pagination(&self) -> &$crate::dto::PaginationParams {
851                &self.pagination
852            }
853        }
854    };
855}
856
857#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash, DeepSizeOf)]
858#[serde(deny_unknown_fields)]
859pub struct PaginationResponse {
860    pub page: i64,
861    pub page_size: i64,
862    /// The total number of items available across all pages of results
863    pub total: i64,
864}
865
866/// Current pagination information
867impl PaginationResponse {
868    pub fn new(page: i64, page_size: i64, total: i64) -> Self {
869        Self { page, page_size, total }
870    }
871
872    pub fn total_pages(&self) -> i64 {
873        // ceil(total / page_size)
874        (self.total + self.page_size - 1) / self.page_size
875    }
876}
877
878#[derive(
879    Clone, Serialize, Debug, Default, Deserialize, PartialEq, ToSchema, Eq, Hash, DeepSizeOf,
880)]
881#[serde(deny_unknown_fields)]
882pub struct StateRequestBody {
883    /// Filters response by contract addresses
884    #[serde(alias = "contractIds")]
885    #[schema(value_type=Option<Vec<String>>)]
886    pub contract_ids: Option<Vec<Bytes>>,
887    /// Does not filter response, only required to correctly apply unconfirmed state
888    /// from ReorgBuffers
889    #[serde(alias = "protocolSystem", default)]
890    pub protocol_system: String,
891    #[serde(default = "VersionParam::default")]
892    pub version: VersionParam,
893    #[serde(default)]
894    pub chain: Chain,
895    #[serde(default)]
896    pub pagination: PaginationParams,
897}
898
899// When INCREASING these limits, please read the warning in the macro definition.
900impl_pagination_limits!(StateRequestBody, compressed = 1200, uncompressed = 100);
901
902impl StateRequestBody {
903    pub fn new(
904        contract_ids: Option<Vec<Bytes>>,
905        protocol_system: String,
906        version: VersionParam,
907        chain: Chain,
908        pagination: PaginationParams,
909    ) -> Self {
910        Self { contract_ids, protocol_system, version, chain, pagination }
911    }
912
913    pub fn from_block(protocol_system: &str, block: BlockParam) -> Self {
914        Self {
915            contract_ids: None,
916            protocol_system: protocol_system.to_string(),
917            version: VersionParam { timestamp: None, block: Some(block.clone()) },
918            chain: block.chain.unwrap_or_default(),
919            pagination: PaginationParams::default(),
920        }
921    }
922
923    pub fn from_timestamp(protocol_system: &str, timestamp: NaiveDateTime, chain: Chain) -> Self {
924        Self {
925            contract_ids: None,
926            protocol_system: protocol_system.to_string(),
927            version: VersionParam { timestamp: Some(timestamp), block: None },
928            chain,
929            pagination: PaginationParams::default(),
930        }
931    }
932}
933
934/// Response from Tycho server for a contract state request.
935#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, DeepSizeOf)]
936pub struct StateRequestResponse {
937    pub accounts: Vec<ResponseAccount>,
938    pub pagination: PaginationResponse,
939}
940
941impl StateRequestResponse {
942    pub fn new(accounts: Vec<ResponseAccount>, pagination: PaginationResponse) -> Self {
943        Self { accounts, pagination }
944    }
945}
946
947#[derive(PartialEq, Clone, Serialize, Deserialize, Default, ToSchema, DeepSizeOf)]
948#[serde(rename = "Account")]
949/// Account struct for the response from Tycho server for a contract state request.
950///
951/// Code is serialized as a hex string instead of a list of bytes.
952pub struct ResponseAccount {
953    pub chain: Chain,
954    /// The address of the account as hex encoded string
955    #[schema(value_type=String, example="0xc9f2e6ea1637E499406986ac50ddC92401ce1f58")]
956    #[serde(with = "hex_bytes")]
957    pub address: Bytes,
958    /// The title of the account usualy specifying its function within the protocol
959    #[schema(value_type=String, example="Protocol Vault")]
960    pub title: String,
961    /// Contract storage map of hex encoded string values
962    #[schema(value_type=HashMap<String, String>, example=json!({"0x....": "0x...."}))]
963    #[serde(with = "hex_hashmap_key_value")]
964    pub slots: HashMap<Bytes, Bytes>,
965    /// The balance of the account in the native token
966    #[schema(value_type=String, example="0x00")]
967    #[serde(with = "hex_bytes")]
968    pub native_balance: Bytes,
969    /// Balances of this account in other tokens (only tokens balance that are
970    /// relevant to the protocol are returned here)
971    #[schema(value_type=HashMap<String, String>, example=json!({"0x....": "0x...."}))]
972    #[serde(with = "hex_hashmap_key_value")]
973    pub token_balances: HashMap<Bytes, Bytes>,
974    /// The accounts code as hex encoded string
975    #[schema(value_type=String, example="0xBADBABE")]
976    #[serde(with = "hex_bytes")]
977    pub code: Bytes,
978    /// The hash of above code
979    #[schema(value_type=String, example="0x123456789")]
980    #[serde(with = "hex_bytes")]
981    pub code_hash: Bytes,
982    /// Transaction hash which last modified native balance
983    #[schema(value_type=String, example="0x8f1133bfb054a23aedfe5d25b1d81b96195396d8b88bd5d4bcf865fc1ae2c3f4")]
984    #[serde(with = "hex_bytes")]
985    pub balance_modify_tx: Bytes,
986    /// Transaction hash which last modified code
987    #[schema(value_type=String, example="0x8f1133bfb054a23aedfe5d25b1d81b96195396d8b88bd5d4bcf865fc1ae2c3f4")]
988    #[serde(with = "hex_bytes")]
989    pub code_modify_tx: Bytes,
990    /// Transaction hash which created the account
991    #[deprecated(note = "The `creation_tx` field is deprecated.")]
992    #[schema(value_type=Option<String>, example="0x8f1133bfb054a23aedfe5d25b1d81b96195396d8b88bd5d4bcf865fc1ae2c3f4")]
993    #[serde(with = "hex_bytes_option")]
994    pub creation_tx: Option<Bytes>,
995}
996
997impl ResponseAccount {
998    #[allow(clippy::too_many_arguments)]
999    pub fn new(
1000        chain: Chain,
1001        address: Bytes,
1002        title: String,
1003        slots: HashMap<Bytes, Bytes>,
1004        native_balance: Bytes,
1005        token_balances: HashMap<Bytes, Bytes>,
1006        code: Bytes,
1007        code_hash: Bytes,
1008        balance_modify_tx: Bytes,
1009        code_modify_tx: Bytes,
1010        creation_tx: Option<Bytes>,
1011    ) -> Self {
1012        Self {
1013            chain,
1014            address,
1015            title,
1016            slots,
1017            native_balance,
1018            token_balances,
1019            code,
1020            code_hash,
1021            balance_modify_tx,
1022            code_modify_tx,
1023            creation_tx,
1024        }
1025    }
1026}
1027
1028/// Implement Debug for ResponseAccount manually to avoid printing the code field.
1029impl fmt::Debug for ResponseAccount {
1030    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1031        f.debug_struct("ResponseAccount")
1032            .field("chain", &self.chain)
1033            .field("address", &self.address)
1034            .field("title", &self.title)
1035            .field("slots", &self.slots)
1036            .field("native_balance", &self.native_balance)
1037            .field("token_balances", &self.token_balances)
1038            .field("code", &format!("[{} bytes]", self.code.len()))
1039            .field("code_hash", &self.code_hash)
1040            .field("balance_modify_tx", &self.balance_modify_tx)
1041            .field("code_modify_tx", &self.code_modify_tx)
1042            .field("creation_tx", &self.creation_tx)
1043            .finish()
1044    }
1045}
1046
1047#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
1048pub struct AccountBalance {
1049    #[serde(with = "hex_bytes")]
1050    pub account: Bytes,
1051    #[serde(with = "hex_bytes")]
1052    pub token: Bytes,
1053    #[serde(with = "hex_bytes")]
1054    pub balance: Bytes,
1055    #[serde(with = "hex_bytes")]
1056    pub modify_tx: Bytes,
1057}
1058
1059#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema)]
1060#[serde(deny_unknown_fields)]
1061pub struct ContractId {
1062    #[serde(with = "hex_bytes")]
1063    #[schema(value_type=String)]
1064    pub address: Bytes,
1065    pub chain: Chain,
1066}
1067
1068/// Uniquely identifies a contract on a specific chain.
1069impl ContractId {
1070    pub fn new(chain: Chain, address: Bytes) -> Self {
1071        Self { address, chain }
1072    }
1073
1074    pub fn address(&self) -> &Bytes {
1075        &self.address
1076    }
1077}
1078
1079impl fmt::Display for ContractId {
1080    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1081        write!(f, "{:?}: 0x{}", self.chain, hex::encode(&self.address))
1082    }
1083}
1084
1085/// The version of the requested state, given as either a timestamp or a block.
1086///
1087/// If block is provided, the state at that exact block is returned. Will error if the block
1088/// has not been processed yet. If timestamp is provided, the state at the latest block before
1089/// that timestamp is returned.
1090/// Defaults to the current time.
1091#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash)]
1092#[serde(deny_unknown_fields)]
1093pub struct VersionParam {
1094    pub timestamp: Option<NaiveDateTime>,
1095    pub block: Option<BlockParam>,
1096}
1097
1098impl DeepSizeOf for VersionParam {
1099    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
1100        if let Some(block) = &self.block {
1101            return block.deep_size_of_children(ctx);
1102        }
1103
1104        0
1105    }
1106}
1107
1108impl VersionParam {
1109    pub fn new(timestamp: Option<NaiveDateTime>, block: Option<BlockParam>) -> Self {
1110        Self { timestamp, block }
1111    }
1112}
1113
1114impl Default for VersionParam {
1115    fn default() -> Self {
1116        VersionParam { timestamp: Some(Utc::now().naive_utc()), block: None }
1117    }
1118}
1119
1120#[deprecated(note = "Use StateRequestBody instead")]
1121#[derive(Serialize, Deserialize, Default, Debug, IntoParams)]
1122pub struct StateRequestParameters {
1123    /// The minimum TVL of the protocol components to return, denoted in the chain's native token.
1124    #[param(default = 0)]
1125    pub tvl_gt: Option<u64>,
1126    /// The minimum inertia of the protocol components to return.
1127    #[param(default = 0)]
1128    pub inertia_min_gt: Option<u64>,
1129    /// Whether to include ERC20 balances in the response.
1130    #[serde(default = "default_include_balances_flag")]
1131    pub include_balances: bool,
1132    #[serde(default)]
1133    pub pagination: PaginationParams,
1134}
1135
1136impl StateRequestParameters {
1137    pub fn new(include_balances: bool) -> Self {
1138        Self {
1139            tvl_gt: None,
1140            inertia_min_gt: None,
1141            include_balances,
1142            pagination: PaginationParams::default(),
1143        }
1144    }
1145
1146    pub fn to_query_string(&self) -> String {
1147        let mut parts = vec![format!("include_balances={}", self.include_balances)];
1148
1149        if let Some(tvl_gt) = self.tvl_gt {
1150            parts.push(format!("tvl_gt={tvl_gt}"));
1151        }
1152
1153        if let Some(inertia) = self.inertia_min_gt {
1154            parts.push(format!("inertia_min_gt={inertia}"));
1155        }
1156
1157        let mut res = parts.join("&");
1158        if !res.is_empty() {
1159            res = format!("?{res}");
1160        }
1161        res
1162    }
1163}
1164
1165#[derive(
1166    Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone, DeepSizeOf,
1167)]
1168#[serde(deny_unknown_fields)]
1169pub struct TokensRequestBody {
1170    /// Filters tokens by addresses
1171    #[serde(alias = "tokenAddresses")]
1172    #[schema(value_type=Option<Vec<String>>)]
1173    pub token_addresses: Option<Vec<Bytes>>,
1174    /// Quality is between 0-100, where:
1175    ///  - 100: Normal ERC-20 Token behavior
1176    ///  - 75: Rebasing token
1177    ///  - 50: Fee-on-transfer token
1178    ///  - 10: Token analysis failed at first detection
1179    ///  - 5: Token analysis failed multiple times (after creation)
1180    ///  - 0: Failed to extract attributes, like Decimal or Symbol
1181    #[serde(default)]
1182    pub min_quality: Option<i32>,
1183    /// Filters tokens by recent trade activity
1184    #[serde(default)]
1185    pub traded_n_days_ago: Option<u64>,
1186    /// Max page size supported is 3000
1187    #[serde(default)]
1188    pub pagination: PaginationParams,
1189    /// Filter tokens by blockchain, default 'ethereum'
1190    #[serde(default)]
1191    pub chain: Chain,
1192}
1193
1194// When INCREASING these limits, please read the warning in the macro definition.
1195impl_pagination_limits!(TokensRequestBody, compressed = 12900, uncompressed = 3000);
1196
1197/// Response from Tycho server for a tokens request.
1198#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash, DeepSizeOf)]
1199pub struct TokensRequestResponse {
1200    pub tokens: Vec<ResponseToken>,
1201    pub pagination: PaginationResponse,
1202}
1203
1204impl TokensRequestResponse {
1205    pub fn new(tokens: Vec<ResponseToken>, pagination_request: &PaginationResponse) -> Self {
1206        Self { tokens, pagination: pagination_request.clone() }
1207    }
1208}
1209
1210#[derive(
1211    PartialEq, Debug, Clone, Serialize, Deserialize, Default, ToSchema, Eq, Hash, DeepSizeOf,
1212)]
1213#[serde(rename = "Token")]
1214/// Token struct for the response from Tycho server for a tokens request.
1215pub struct ResponseToken {
1216    pub chain: Chain,
1217    /// The address of this token as hex encoded string
1218    #[schema(value_type=String, example="0xc9f2e6ea1637E499406986ac50ddC92401ce1f58")]
1219    #[serde(with = "hex_bytes")]
1220    pub address: Bytes,
1221    /// A shorthand symbol for this token (not unique)
1222    #[schema(value_type=String, example="WETH")]
1223    pub symbol: String,
1224    /// The number of decimals used to represent token values
1225    pub decimals: u32,
1226    /// The tax this token charges on transfers in basis points
1227    pub tax: u64,
1228    /// Gas usage of the token, currently is always a single averaged value
1229    pub gas: Vec<Option<u64>>,
1230    /// Quality is between 0-100, where:
1231    ///  - 100: Normal ERC-20 Token behavior
1232    ///  - 75: Rebasing token
1233    ///  - 50: Fee-on-transfer token
1234    ///  - 10: Token analysis failed at first detection
1235    ///  - 5: Token analysis failed multiple times (after creation)
1236    ///  - 0: Failed to extract attributes, like Decimal or Symbol
1237    pub quality: u32,
1238}
1239
1240impl From<models::token::Token> for ResponseToken {
1241    fn from(value: models::token::Token) -> Self {
1242        Self {
1243            chain: value.chain.into(),
1244            address: value.address,
1245            symbol: value.symbol,
1246            decimals: value.decimals,
1247            tax: value.tax,
1248            gas: value.gas,
1249            quality: value.quality,
1250        }
1251    }
1252}
1253
1254#[derive(Serialize, Deserialize, Debug, Default, ToSchema, Clone, DeepSizeOf)]
1255#[serde(deny_unknown_fields)]
1256pub struct ProtocolComponentsRequestBody {
1257    /// Filters by protocol, required to correctly apply unconfirmed state from
1258    /// ReorgBuffers
1259    pub protocol_system: String,
1260    /// Filter by component ids
1261    #[schema(value_type=Option<Vec<String>>)]
1262    #[serde(alias = "componentAddresses")]
1263    pub component_ids: Option<Vec<ComponentId>>,
1264    /// The minimum TVL of the protocol components to return, denoted in the chain's
1265    /// native token.
1266    #[serde(default)]
1267    pub tvl_gt: Option<f64>,
1268    #[serde(default)]
1269    pub chain: Chain,
1270    /// Max page size supported is 500
1271    #[serde(default)]
1272    pub pagination: PaginationParams,
1273}
1274
1275// When INCREASING these limits, please read the warning in the macro definition.
1276impl_pagination_limits!(ProtocolComponentsRequestBody, compressed = 2550, uncompressed = 500);
1277
1278// Implement PartialEq where tvl is considered equal if the difference is less than 1e-6
1279impl PartialEq for ProtocolComponentsRequestBody {
1280    fn eq(&self, other: &Self) -> bool {
1281        let tvl_close_enough = match (self.tvl_gt, other.tvl_gt) {
1282            (Some(a), Some(b)) => (a - b).abs() < 1e-6,
1283            (None, None) => true,
1284            _ => false,
1285        };
1286
1287        self.protocol_system == other.protocol_system &&
1288            self.component_ids == other.component_ids &&
1289            tvl_close_enough &&
1290            self.chain == other.chain &&
1291            self.pagination == other.pagination
1292    }
1293}
1294
1295// Implement Eq without any new logic
1296impl Eq for ProtocolComponentsRequestBody {}
1297
1298impl Hash for ProtocolComponentsRequestBody {
1299    fn hash<H: Hasher>(&self, state: &mut H) {
1300        self.protocol_system.hash(state);
1301        self.component_ids.hash(state);
1302
1303        // Handle the f64 `tvl_gt` field by converting it into a hashable integer
1304        if let Some(tvl) = self.tvl_gt {
1305            // Convert f64 to bits and hash those bits
1306            tvl.to_bits().hash(state);
1307        } else {
1308            // Use a constant value to represent None
1309            state.write_u8(0);
1310        }
1311
1312        self.chain.hash(state);
1313        self.pagination.hash(state);
1314    }
1315}
1316
1317impl ProtocolComponentsRequestBody {
1318    pub fn system_filtered(system: &str, tvl_gt: Option<f64>, chain: Chain) -> Self {
1319        Self {
1320            protocol_system: system.to_string(),
1321            component_ids: None,
1322            tvl_gt,
1323            chain,
1324            pagination: Default::default(),
1325        }
1326    }
1327
1328    pub fn id_filtered(system: &str, ids: Vec<String>, chain: Chain) -> Self {
1329        Self {
1330            protocol_system: system.to_string(),
1331            component_ids: Some(ids),
1332            tvl_gt: None,
1333            chain,
1334            pagination: Default::default(),
1335        }
1336    }
1337}
1338
1339impl ProtocolComponentsRequestBody {
1340    pub fn new(
1341        protocol_system: String,
1342        component_ids: Option<Vec<String>>,
1343        tvl_gt: Option<f64>,
1344        chain: Chain,
1345        pagination: PaginationParams,
1346    ) -> Self {
1347        Self { protocol_system, component_ids, tvl_gt, chain, pagination }
1348    }
1349}
1350
1351#[deprecated(note = "Use ProtocolComponentsRequestBody instead")]
1352#[derive(Serialize, Deserialize, Default, Debug, IntoParams)]
1353pub struct ProtocolComponentRequestParameters {
1354    /// The minimum TVL of the protocol components to return, denoted in the chain's native token.
1355    #[param(default = 0)]
1356    pub tvl_gt: Option<f64>,
1357}
1358
1359impl ProtocolComponentRequestParameters {
1360    pub fn tvl_filtered(min_tvl: f64) -> Self {
1361        Self { tvl_gt: Some(min_tvl) }
1362    }
1363}
1364
1365impl ProtocolComponentRequestParameters {
1366    pub fn to_query_string(&self) -> String {
1367        if let Some(tvl_gt) = self.tvl_gt {
1368            return format!("?tvl_gt={tvl_gt}");
1369        }
1370        String::new()
1371    }
1372}
1373
1374/// Response from Tycho server for a protocol components request.
1375#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, DeepSizeOf)]
1376pub struct ProtocolComponentRequestResponse {
1377    pub protocol_components: Vec<ProtocolComponent>,
1378    pub pagination: PaginationResponse,
1379}
1380
1381impl ProtocolComponentRequestResponse {
1382    pub fn new(
1383        protocol_components: Vec<ProtocolComponent>,
1384        pagination: PaginationResponse,
1385    ) -> Self {
1386        Self { protocol_components, pagination }
1387    }
1388}
1389
1390#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash)]
1391#[serde(deny_unknown_fields)]
1392#[deprecated]
1393pub struct ProtocolId {
1394    pub id: String,
1395    pub chain: Chain,
1396}
1397
1398impl From<ProtocolId> for String {
1399    fn from(protocol_id: ProtocolId) -> Self {
1400        protocol_id.id
1401    }
1402}
1403
1404impl AsRef<str> for ProtocolId {
1405    fn as_ref(&self) -> &str {
1406        &self.id
1407    }
1408}
1409
1410/// Protocol State struct for the response from Tycho server for a protocol state request.
1411#[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize, ToSchema, DeepSizeOf)]
1412pub struct ResponseProtocolState {
1413    /// Component id this state belongs to
1414    pub component_id: String,
1415    /// Attributes of the component. If an attribute's value is a `bigint`,
1416    /// it will be encoded as a big endian signed hex string.
1417    #[schema(value_type=HashMap<String, String>)]
1418    #[serde(with = "hex_hashmap_value")]
1419    pub attributes: HashMap<String, Bytes>,
1420    /// Sum aggregated balances of the component
1421    #[schema(value_type=HashMap<String, String>)]
1422    #[serde(with = "hex_hashmap_key_value")]
1423    pub balances: HashMap<Bytes, Bytes>,
1424}
1425
1426impl From<models::protocol::ProtocolComponentState> for ResponseProtocolState {
1427    fn from(value: models::protocol::ProtocolComponentState) -> Self {
1428        Self {
1429            component_id: value.component_id,
1430            attributes: value.attributes,
1431            balances: value.balances,
1432        }
1433    }
1434}
1435
1436fn default_include_balances_flag() -> bool {
1437    true
1438}
1439
1440/// Max page size supported is 100
1441#[derive(Clone, Debug, Serialize, PartialEq, ToSchema, Default, Eq, Hash, DeepSizeOf)]
1442#[serde(deny_unknown_fields)]
1443pub struct ProtocolStateRequestBody {
1444    /// Filters response by protocol components ids
1445    #[serde(alias = "protocolIds")]
1446    pub protocol_ids: Option<Vec<String>>,
1447    /// Filters by protocol, required to correctly apply unconfirmed state from
1448    /// ReorgBuffers
1449    #[serde(alias = "protocolSystem")]
1450    pub protocol_system: String,
1451    #[serde(default)]
1452    pub chain: Chain,
1453    /// Whether to include account balances in the response. Defaults to true.
1454    #[serde(default = "default_include_balances_flag")]
1455    pub include_balances: bool,
1456    #[serde(default = "VersionParam::default")]
1457    pub version: VersionParam,
1458    #[serde(default)]
1459    pub pagination: PaginationParams,
1460}
1461
1462// When INCREASING these limits, please read the warning in the macro definition.
1463impl_pagination_limits!(ProtocolStateRequestBody, compressed = 360, uncompressed = 100);
1464
1465impl ProtocolStateRequestBody {
1466    pub fn id_filtered<I, T>(ids: I) -> Self
1467    where
1468        I: IntoIterator<Item = T>,
1469        T: Into<String>,
1470    {
1471        Self {
1472            protocol_ids: Some(
1473                ids.into_iter()
1474                    .map(Into::into)
1475                    .collect(),
1476            ),
1477            ..Default::default()
1478        }
1479    }
1480}
1481
1482/// Custom deserializer for ProtocolStateRequestBody to support backwards compatibility with the old
1483/// ProtocolIds format.
1484/// To be removed when the old format is no longer supported.
1485impl<'de> Deserialize<'de> for ProtocolStateRequestBody {
1486    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1487    where
1488        D: Deserializer<'de>,
1489    {
1490        #[derive(Deserialize)]
1491        #[serde(untagged)]
1492        enum ProtocolIdOrString {
1493            Old(Vec<ProtocolId>),
1494            New(Vec<String>),
1495        }
1496
1497        struct ProtocolStateRequestBodyVisitor;
1498
1499        impl<'de> de::Visitor<'de> for ProtocolStateRequestBodyVisitor {
1500            type Value = ProtocolStateRequestBody;
1501
1502            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
1503                formatter.write_str("struct ProtocolStateRequestBody")
1504            }
1505
1506            fn visit_map<V>(self, mut map: V) -> Result<ProtocolStateRequestBody, V::Error>
1507            where
1508                V: de::MapAccess<'de>,
1509            {
1510                let mut protocol_ids = None;
1511                let mut protocol_system = None;
1512                let mut version = None;
1513                let mut chain = None;
1514                let mut include_balances = None;
1515                let mut pagination = None;
1516
1517                while let Some(key) = map.next_key::<String>()? {
1518                    match key.as_str() {
1519                        "protocol_ids" | "protocolIds" => {
1520                            let value: ProtocolIdOrString = map.next_value()?;
1521                            protocol_ids = match value {
1522                                ProtocolIdOrString::Old(ids) => {
1523                                    Some(ids.into_iter().map(|p| p.id).collect())
1524                                }
1525                                ProtocolIdOrString::New(ids_str) => Some(ids_str),
1526                            };
1527                        }
1528                        "protocol_system" | "protocolSystem" => {
1529                            protocol_system = Some(map.next_value()?);
1530                        }
1531                        "version" => {
1532                            version = Some(map.next_value()?);
1533                        }
1534                        "chain" => {
1535                            chain = Some(map.next_value()?);
1536                        }
1537                        "include_balances" => {
1538                            include_balances = Some(map.next_value()?);
1539                        }
1540                        "pagination" => {
1541                            pagination = Some(map.next_value()?);
1542                        }
1543                        _ => {
1544                            return Err(de::Error::unknown_field(
1545                                &key,
1546                                &[
1547                                    "contract_ids",
1548                                    "protocol_system",
1549                                    "version",
1550                                    "chain",
1551                                    "include_balances",
1552                                    "pagination",
1553                                ],
1554                            ))
1555                        }
1556                    }
1557                }
1558
1559                Ok(ProtocolStateRequestBody {
1560                    protocol_ids,
1561                    protocol_system: protocol_system.unwrap_or_default(),
1562                    version: version.unwrap_or_else(VersionParam::default),
1563                    chain: chain.unwrap_or_else(Chain::default),
1564                    include_balances: include_balances.unwrap_or(true),
1565                    pagination: pagination.unwrap_or_else(PaginationParams::default),
1566                })
1567            }
1568        }
1569
1570        deserializer.deserialize_struct(
1571            "ProtocolStateRequestBody",
1572            &[
1573                "contract_ids",
1574                "protocol_system",
1575                "version",
1576                "chain",
1577                "include_balances",
1578                "pagination",
1579            ],
1580            ProtocolStateRequestBodyVisitor,
1581        )
1582    }
1583}
1584
1585#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, DeepSizeOf)]
1586pub struct ProtocolStateRequestResponse {
1587    pub states: Vec<ResponseProtocolState>,
1588    pub pagination: PaginationResponse,
1589}
1590
1591impl ProtocolStateRequestResponse {
1592    pub fn new(states: Vec<ResponseProtocolState>, pagination: PaginationResponse) -> Self {
1593        Self { states, pagination }
1594    }
1595}
1596
1597#[derive(Serialize, Clone, PartialEq, Hash, Eq)]
1598pub struct ProtocolComponentId {
1599    pub chain: Chain,
1600    pub system: String,
1601    pub id: String,
1602}
1603
1604#[derive(Debug, Serialize, ToSchema)]
1605#[serde(tag = "status", content = "message")]
1606#[schema(example = json!({"status": "NotReady", "message": "No db connection"}))]
1607pub enum Health {
1608    Ready,
1609    Starting(String),
1610    NotReady(String),
1611}
1612
1613#[derive(Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone)]
1614#[serde(deny_unknown_fields)]
1615pub struct ProtocolSystemsRequestBody {
1616    #[serde(default)]
1617    pub chain: Chain,
1618    #[serde(default)]
1619    pub pagination: PaginationParams,
1620}
1621
1622// When INCREASING these limits, please read the warning in the macro definition.
1623impl_pagination_limits!(ProtocolSystemsRequestBody, compressed = 100, uncompressed = 100);
1624
1625#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash)]
1626pub struct ProtocolSystemsRequestResponse {
1627    /// List of currently supported protocol systems
1628    pub protocol_systems: Vec<String>,
1629    pub pagination: PaginationResponse,
1630}
1631
1632impl ProtocolSystemsRequestResponse {
1633    pub fn new(protocol_systems: Vec<String>, pagination: PaginationResponse) -> Self {
1634        Self { protocol_systems, pagination }
1635    }
1636}
1637
1638#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
1639pub struct DCIUpdate {
1640    /// Map of component id to the new entrypoints associated with the component
1641    pub new_entrypoints: HashMap<ComponentId, HashSet<EntryPoint>>,
1642    /// Map of entrypoint id to the new entrypoint params associtated with it (and optionally the
1643    /// component linked to those params)
1644    pub new_entrypoint_params: HashMap<String, HashSet<(TracingParams, String)>>,
1645    /// Map of entrypoint id to its trace result
1646    pub trace_results: HashMap<String, TracingResult>,
1647}
1648
1649impl From<models::blockchain::DCIUpdate> for DCIUpdate {
1650    fn from(value: models::blockchain::DCIUpdate) -> Self {
1651        Self {
1652            new_entrypoints: value
1653                .new_entrypoints
1654                .into_iter()
1655                .map(|(k, v)| {
1656                    (
1657                        k,
1658                        v.into_iter()
1659                            .map(|v| v.into())
1660                            .collect(),
1661                    )
1662                })
1663                .collect(),
1664            new_entrypoint_params: value
1665                .new_entrypoint_params
1666                .into_iter()
1667                .map(|(k, v)| {
1668                    (
1669                        k,
1670                        v.into_iter()
1671                            .map(|(params, i)| (params.into(), i))
1672                            .collect(),
1673                    )
1674                })
1675                .collect(),
1676            trace_results: value
1677                .trace_results
1678                .into_iter()
1679                .map(|(k, v)| (k, v.into()))
1680                .collect(),
1681        }
1682    }
1683}
1684
1685#[derive(Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone)]
1686#[serde(deny_unknown_fields)]
1687pub struct ComponentTvlRequestBody {
1688    #[serde(default)]
1689    pub chain: Chain,
1690    /// Filters protocol components by protocol system
1691    /// Useful when `component_ids` is omitted to fetch all components under a specific system.
1692    #[serde(alias = "protocolSystem")]
1693    pub protocol_system: Option<String>,
1694    #[serde(default)]
1695    pub component_ids: Option<Vec<String>>,
1696    #[serde(default)]
1697    pub pagination: PaginationParams,
1698}
1699
1700// When INCREASING these limits, please read the warning in the macro definition.
1701impl_pagination_limits!(ComponentTvlRequestBody, compressed = 100, uncompressed = 100);
1702
1703impl ComponentTvlRequestBody {
1704    pub fn system_filtered(system: &str, chain: Chain) -> Self {
1705        Self {
1706            chain,
1707            protocol_system: Some(system.to_string()),
1708            component_ids: None,
1709            pagination: Default::default(),
1710        }
1711    }
1712
1713    pub fn id_filtered(ids: Vec<String>, chain: Chain) -> Self {
1714        Self {
1715            chain,
1716            protocol_system: None,
1717            component_ids: Some(ids),
1718            pagination: Default::default(),
1719        }
1720    }
1721}
1722// #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash)]
1723#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema)]
1724pub struct ComponentTvlRequestResponse {
1725    pub tvl: HashMap<String, f64>,
1726    pub pagination: PaginationResponse,
1727}
1728
1729impl ComponentTvlRequestResponse {
1730    pub fn new(tvl: HashMap<String, f64>, pagination: PaginationResponse) -> Self {
1731        Self { tvl, pagination }
1732    }
1733}
1734
1735#[derive(
1736    Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone, DeepSizeOf,
1737)]
1738pub struct TracedEntryPointRequestBody {
1739    #[serde(default)]
1740    pub chain: Chain,
1741    /// Filters by protocol, required to correctly apply unconfirmed state from
1742    /// ReorgBuffers
1743    pub protocol_system: String,
1744    /// Filter by component ids
1745    #[schema(value_type = Option<Vec<String>>)]
1746    pub component_ids: Option<Vec<ComponentId>>,
1747    /// Max page size supported is 100
1748    #[serde(default)]
1749    pub pagination: PaginationParams,
1750}
1751
1752// When INCREASING these limits, please read the warning in the macro definition.
1753impl_pagination_limits!(TracedEntryPointRequestBody, compressed = 100, uncompressed = 100);
1754
1755#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash, DeepSizeOf)]
1756pub struct EntryPoint {
1757    #[schema(example = "0xEdf63cce4bA70cbE74064b7687882E71ebB0e988:getRate()")]
1758    /// Entry point id.
1759    pub external_id: String,
1760    #[schema(value_type=String, example="0x8f4E8439b970363648421C692dd897Fb9c0Bd1D9")]
1761    #[serde(with = "hex_bytes")]
1762    /// The address of the contract to trace.
1763    pub target: Bytes,
1764    #[schema(example = "getRate()")]
1765    /// The signature of the function to trace.
1766    pub signature: String,
1767}
1768
1769#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema, Eq, Hash, DeepSizeOf)]
1770pub enum StorageOverride {
1771    /// Applies changes incrementally to the existing account storage.
1772    /// Only modifies the specific storage slots provided in the map while
1773    /// preserving all other storage slots.
1774    #[schema(value_type=HashMap<String, String>)]
1775    Diff(BTreeMap<StoreKey, StoreVal>),
1776
1777    /// Completely replaces the account's storage state.
1778    /// Only the storage slots provided in the map will exist after the operation,
1779    /// and any existing storage slots not included will be cleared/zeroed.
1780    #[schema(value_type=HashMap<String, String>)]
1781    Replace(BTreeMap<StoreKey, StoreVal>),
1782}
1783
1784impl From<models::blockchain::StorageOverride> for StorageOverride {
1785    fn from(value: models::blockchain::StorageOverride) -> Self {
1786        match value {
1787            models::blockchain::StorageOverride::Diff(diff) => StorageOverride::Diff(diff),
1788            models::blockchain::StorageOverride::Replace(replace) => {
1789                StorageOverride::Replace(replace)
1790            }
1791        }
1792    }
1793}
1794
1795/// State overrides for an account.
1796///
1797/// Used to modify account state. Commonly used for testing contract interactions with specific
1798/// state conditions or simulating transactions with modified balances/code.
1799#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema, Eq, Hash, DeepSizeOf)]
1800pub struct AccountOverrides {
1801    /// Storage slots to override
1802    pub slots: Option<StorageOverride>,
1803    #[schema(value_type=Option<String>)]
1804    /// Native token balance override
1805    pub native_balance: Option<Balance>,
1806    #[schema(value_type=Option<String>)]
1807    /// Contract code override
1808    pub code: Option<Code>,
1809}
1810
1811impl From<models::blockchain::AccountOverrides> for AccountOverrides {
1812    fn from(value: models::blockchain::AccountOverrides) -> Self {
1813        AccountOverrides {
1814            slots: value.slots.map(|s| s.into()),
1815            native_balance: value.native_balance,
1816            code: value.code,
1817        }
1818    }
1819}
1820
1821#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash, DeepSizeOf)]
1822pub struct RPCTracerParams {
1823    /// The caller address of the transaction, if not provided tracing uses the default value
1824    /// for an address defined by the VM.
1825    #[schema(value_type=Option<String>)]
1826    #[serde(with = "hex_bytes_option", default)]
1827    pub caller: Option<Bytes>,
1828    /// The call data used for the tracing call, this needs to include the function selector
1829    #[schema(value_type=String, example="0x679aefce")]
1830    #[serde(with = "hex_bytes")]
1831    pub calldata: Bytes,
1832    /// Optionally allow for state overrides so that the call works as expected
1833    pub state_overrides: Option<BTreeMap<Address, AccountOverrides>>,
1834    /// Addresses to prune from trace results. Useful for hooks that use mock
1835    /// accounts/routers that shouldn't be tracked in the final DCI results.
1836    #[schema(value_type=Option<Vec<String>>)]
1837    #[serde(default)]
1838    pub prune_addresses: Option<Vec<Address>>,
1839}
1840
1841impl From<models::blockchain::RPCTracerParams> for RPCTracerParams {
1842    fn from(value: models::blockchain::RPCTracerParams) -> Self {
1843        RPCTracerParams {
1844            caller: value.caller,
1845            calldata: value.calldata,
1846            state_overrides: value.state_overrides.map(|overrides| {
1847                overrides
1848                    .into_iter()
1849                    .map(|(address, account_overrides)| (address, account_overrides.into()))
1850                    .collect()
1851            }),
1852            prune_addresses: value.prune_addresses,
1853        }
1854    }
1855}
1856
1857#[derive(Deserialize, Serialize, Debug, PartialEq, Eq, Clone, Hash, DeepSizeOf, ToSchema)]
1858#[serde(tag = "method", rename_all = "lowercase")]
1859pub enum TracingParams {
1860    /// Uses RPC calls to retrieve the called addresses and retriggers
1861    RPCTracer(RPCTracerParams),
1862}
1863
1864impl From<models::blockchain::TracingParams> for TracingParams {
1865    fn from(value: models::blockchain::TracingParams) -> Self {
1866        match value {
1867            models::blockchain::TracingParams::RPCTracer(params) => {
1868                TracingParams::RPCTracer(params.into())
1869            }
1870        }
1871    }
1872}
1873
1874impl From<models::blockchain::EntryPoint> for EntryPoint {
1875    fn from(value: models::blockchain::EntryPoint) -> Self {
1876        Self { external_id: value.external_id, target: value.target, signature: value.signature }
1877    }
1878}
1879
1880#[derive(Serialize, Deserialize, Debug, PartialEq, ToSchema, Eq, Clone, DeepSizeOf)]
1881pub struct EntryPointWithTracingParams {
1882    /// The entry point object
1883    pub entry_point: EntryPoint,
1884    /// The parameters used
1885    pub params: TracingParams,
1886}
1887
1888impl From<models::blockchain::EntryPointWithTracingParams> for EntryPointWithTracingParams {
1889    fn from(value: models::blockchain::EntryPointWithTracingParams) -> Self {
1890        Self { entry_point: value.entry_point.into(), params: value.params.into() }
1891    }
1892}
1893
1894#[derive(
1895    Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Serialize, Deserialize, DeepSizeOf,
1896)]
1897pub struct AddressStorageLocation {
1898    pub key: StoreKey,
1899    pub offset: u8,
1900}
1901
1902impl AddressStorageLocation {
1903    pub fn new(key: StoreKey, offset: u8) -> Self {
1904        Self { key, offset }
1905    }
1906}
1907
1908impl From<models::blockchain::AddressStorageLocation> for AddressStorageLocation {
1909    fn from(value: models::blockchain::AddressStorageLocation) -> Self {
1910        Self { key: value.key, offset: value.offset }
1911    }
1912}
1913
1914fn deserialize_retriggers_from_value(
1915    value: &serde_json::Value,
1916) -> Result<HashSet<(StoreKey, AddressStorageLocation)>, String> {
1917    use serde::Deserialize;
1918    use serde_json::Value;
1919
1920    let mut result = HashSet::new();
1921
1922    if let Value::Array(items) = value {
1923        for item in items {
1924            if let Value::Array(pair) = item {
1925                if pair.len() == 2 {
1926                    let key = StoreKey::deserialize(&pair[0])
1927                        .map_err(|e| format!("Failed to deserialize key: {}", e))?;
1928
1929                    // Handle both old format (string) and new format (AddressStorageLocation)
1930                    let addr_storage = match &pair[1] {
1931                        Value::String(_) => {
1932                            // Old format: just a string key with offset defaulted to 0
1933                            let storage_key = StoreKey::deserialize(&pair[1]).map_err(|e| {
1934                                format!("Failed to deserialize old format storage key: {}", e)
1935                            })?;
1936                            AddressStorageLocation::new(storage_key, 12)
1937                        }
1938                        Value::Object(_) => {
1939                            // New format: AddressStorageLocation struct
1940                            AddressStorageLocation::deserialize(&pair[1]).map_err(|e| {
1941                                format!("Failed to deserialize AddressStorageLocation: {}", e)
1942                            })?
1943                        }
1944                        _ => return Err("Invalid retrigger format".to_string()),
1945                    };
1946
1947                    result.insert((key, addr_storage));
1948                }
1949            }
1950        }
1951    }
1952
1953    Ok(result)
1954}
1955
1956#[derive(Serialize, Debug, Default, PartialEq, ToSchema, Eq, Clone, DeepSizeOf)]
1957pub struct TracingResult {
1958    #[schema(value_type=HashSet<(String, String)>)]
1959    pub retriggers: HashSet<(StoreKey, AddressStorageLocation)>,
1960    #[schema(value_type=HashMap<String,HashSet<String>>)]
1961    pub accessed_slots: HashMap<Address, HashSet<StoreKey>>,
1962}
1963
1964/// Deserialize TracingResult with backward compatibility for retriggers
1965/// TODO: remove this after offset detection is deployed in production
1966impl<'de> Deserialize<'de> for TracingResult {
1967    fn deserialize<D>(deserializer: D) -> Result<TracingResult, D::Error>
1968    where
1969        D: Deserializer<'de>,
1970    {
1971        use serde::de::Error;
1972        use serde_json::Value;
1973
1974        let value = Value::deserialize(deserializer)?;
1975        let mut result = TracingResult::default();
1976
1977        if let Value::Object(map) = value {
1978            // Deserialize retriggers using our custom deserializer
1979            if let Some(retriggers_value) = map.get("retriggers") {
1980                result.retriggers =
1981                    deserialize_retriggers_from_value(retriggers_value).map_err(|e| {
1982                        D::Error::custom(format!("Failed to deserialize retriggers: {}", e))
1983                    })?;
1984            }
1985
1986            // Deserialize accessed_slots normally
1987            if let Some(accessed_slots_value) = map.get("accessed_slots") {
1988                result.accessed_slots = serde_json::from_value(accessed_slots_value.clone())
1989                    .map_err(|e| {
1990                        D::Error::custom(format!("Failed to deserialize accessed_slots: {}", e))
1991                    })?;
1992            }
1993        }
1994
1995        Ok(result)
1996    }
1997}
1998
1999impl From<models::blockchain::TracingResult> for TracingResult {
2000    fn from(value: models::blockchain::TracingResult) -> Self {
2001        TracingResult {
2002            retriggers: value
2003                .retriggers
2004                .into_iter()
2005                .map(|(k, v)| (k, v.into()))
2006                .collect(),
2007            accessed_slots: value.accessed_slots,
2008        }
2009    }
2010}
2011
2012#[derive(Serialize, PartialEq, ToSchema, Eq, Clone, Debug, Deserialize, DeepSizeOf)]
2013pub struct TracedEntryPointRequestResponse {
2014    /// Map of protocol component id to a list of a tuple containing each entry point with its
2015    /// tracing parameters and its corresponding tracing results.
2016    #[schema(value_type = HashMap<String, Vec<(EntryPointWithTracingParams, TracingResult)>>)]
2017    pub traced_entry_points:
2018        HashMap<ComponentId, Vec<(EntryPointWithTracingParams, TracingResult)>>,
2019    pub pagination: PaginationResponse,
2020}
2021impl From<TracedEntryPointRequestResponse> for DCIUpdate {
2022    fn from(response: TracedEntryPointRequestResponse) -> Self {
2023        let mut new_entrypoints = HashMap::new();
2024        let mut new_entrypoint_params = HashMap::new();
2025        let mut trace_results = HashMap::new();
2026
2027        for (component, traces) in response.traced_entry_points {
2028            let mut entrypoints = HashSet::new();
2029
2030            for (entrypoint, trace) in traces {
2031                let entrypoint_id = entrypoint
2032                    .entry_point
2033                    .external_id
2034                    .clone();
2035
2036                // Collect entrypoints
2037                entrypoints.insert(entrypoint.entry_point.clone());
2038
2039                // Collect entrypoint params
2040                new_entrypoint_params
2041                    .entry(entrypoint_id.clone())
2042                    .or_insert_with(HashSet::new)
2043                    .insert((entrypoint.params, component.clone()));
2044
2045                // Collect trace results
2046                trace_results
2047                    .entry(entrypoint_id)
2048                    .and_modify(|existing_trace: &mut TracingResult| {
2049                        // Merge traces for the same entrypoint
2050                        existing_trace
2051                            .retriggers
2052                            .extend(trace.retriggers.clone());
2053                        for (address, slots) in trace.accessed_slots.clone() {
2054                            existing_trace
2055                                .accessed_slots
2056                                .entry(address)
2057                                .or_default()
2058                                .extend(slots);
2059                        }
2060                    })
2061                    .or_insert(trace);
2062            }
2063
2064            if !entrypoints.is_empty() {
2065                new_entrypoints.insert(component, entrypoints);
2066            }
2067        }
2068
2069        DCIUpdate { new_entrypoints, new_entrypoint_params, trace_results }
2070    }
2071}
2072
2073#[derive(Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Clone)]
2074pub struct AddEntryPointRequestBody {
2075    #[serde(default)]
2076    pub chain: Chain,
2077    #[schema(value_type=String)]
2078    #[serde(default)]
2079    pub block_hash: Bytes,
2080    /// The map of component ids to their tracing params to insert
2081    #[schema(value_type = Vec<(String, Vec<EntryPointWithTracingParams>)>)]
2082    pub entry_points_with_tracing_data: Vec<(ComponentId, Vec<EntryPointWithTracingParams>)>,
2083}
2084
2085#[derive(Serialize, PartialEq, ToSchema, Eq, Clone, Debug, Deserialize)]
2086pub struct AddEntryPointRequestResponse {
2087    /// Map of protocol component id to a list of a tuple containing each entry point with its
2088    /// tracing parameters and its corresponding tracing results.
2089    #[schema(value_type = HashMap<String, Vec<(EntryPointWithTracingParams, TracingResult)>>)]
2090    pub traced_entry_points:
2091        HashMap<ComponentId, Vec<(EntryPointWithTracingParams, TracingResult)>>,
2092}
2093
2094#[cfg(test)]
2095mod test {
2096    use std::str::FromStr;
2097
2098    use maplit::hashmap;
2099    use rstest::rstest;
2100
2101    use super::*;
2102
2103    /// Test backward compatibility for Command::Subscribe compression field.
2104    /// Should default to false when not specified.
2105    #[rstest]
2106    #[case::legacy_format(None, false)]
2107    #[case::explicit_true(Some(true), true)]
2108    #[case::explicit_false(Some(false), false)]
2109    fn test_subscribe_compression_backward_compatibility(
2110        #[case] compression: Option<bool>,
2111        #[case] expected: bool,
2112    ) {
2113        use serde_json::json;
2114
2115        let mut json_value = json!({
2116            "method": "subscribe",
2117            "extractor_id": {
2118                "chain": "ethereum",
2119                "name": "test"
2120            },
2121            "include_state": true
2122        });
2123
2124        if let Some(value) = compression {
2125            json_value["compression"] = json!(value);
2126        }
2127
2128        let command: Command =
2129            serde_json::from_value(json_value).expect("Failed to deserialize Subscribe command");
2130
2131        if let Command::Subscribe { compression, .. } = command {
2132            assert_eq!(compression, expected);
2133        } else {
2134            panic!("Expected Subscribe command");
2135        }
2136    }
2137
2138    /// Test backward compatibility for Command::Subscribe partial_blocks field.
2139    /// Should default to false when not specified.
2140    #[rstest]
2141    #[case::legacy_format(None, false)]
2142    #[case::explicit_true(Some(true), true)]
2143    #[case::explicit_false(Some(false), false)]
2144    fn test_subscribe_partial_blocks_backward_compatibility(
2145        #[case] partial_blocks: Option<bool>,
2146        #[case] expected: bool,
2147    ) {
2148        use serde_json::json;
2149
2150        let mut json_value = json!({
2151            "method": "subscribe",
2152            "extractor_id": {
2153                "chain": "ethereum",
2154                "name": "test"
2155            },
2156            "include_state": true
2157        });
2158
2159        if let Some(value) = partial_blocks {
2160            json_value["partial_blocks"] = json!(value);
2161        }
2162
2163        let command: Command =
2164            serde_json::from_value(json_value).expect("Failed to deserialize Subscribe command");
2165
2166        if let Command::Subscribe { partial_blocks, .. } = command {
2167            assert_eq!(partial_blocks, expected);
2168        } else {
2169            panic!("Expected Subscribe command");
2170        }
2171    }
2172
2173    #[test]
2174    fn test_tracing_result_backward_compatibility() {
2175        use serde_json::json;
2176
2177        // Test old format (string storage locations)
2178        let old_format_json = json!({
2179            "retriggers": [
2180                ["0x01", "0x02"],
2181                ["0x03", "0x04"]
2182            ],
2183            "accessed_slots": {
2184                "0x05": ["0x06", "0x07"]
2185            }
2186        });
2187
2188        let result: TracingResult = serde_json::from_value(old_format_json).unwrap();
2189
2190        // Check that retriggers were deserialized correctly with offset 0
2191        assert_eq!(result.retriggers.len(), 2);
2192        let retriggers_vec: Vec<_> = result.retriggers.iter().collect();
2193        assert!(retriggers_vec.iter().any(|(k, v)| {
2194            k == &Bytes::from("0x01") && v.key == Bytes::from("0x02") && v.offset == 12
2195        }));
2196        assert!(retriggers_vec.iter().any(|(k, v)| {
2197            k == &Bytes::from("0x03") && v.key == Bytes::from("0x04") && v.offset == 12
2198        }));
2199
2200        // Test new format (AddressStorageLocation objects)
2201        let new_format_json = json!({
2202            "retriggers": [
2203                ["0x01", {"key": "0x02", "offset": 12}],
2204                ["0x03", {"key": "0x04", "offset": 5}]
2205            ],
2206            "accessed_slots": {
2207                "0x05": ["0x06", "0x07"]
2208            }
2209        });
2210
2211        let result2: TracingResult = serde_json::from_value(new_format_json).unwrap();
2212
2213        // Check that new format retriggers were deserialized correctly with proper offsets
2214        assert_eq!(result2.retriggers.len(), 2);
2215        let retriggers_vec2: Vec<_> = result2.retriggers.iter().collect();
2216        assert!(retriggers_vec2.iter().any(|(k, v)| {
2217            k == &Bytes::from("0x01") && v.key == Bytes::from("0x02") && v.offset == 12
2218        }));
2219        assert!(retriggers_vec2.iter().any(|(k, v)| {
2220            k == &Bytes::from("0x03") && v.key == Bytes::from("0x04") && v.offset == 5
2221        }));
2222    }
2223
2224    #[rstest]
2225    #[case::legacy_format(None, None)]
2226    #[case::full_block(Some(None), None)]
2227    #[case::partial_block(Some(Some(1)), Some(1))]
2228    fn test_block_changes_is_partial_backward_compatibility(
2229        #[case] has_partial_value: Option<Option<u32>>,
2230        #[case] expected: Option<u32>,
2231    ) {
2232        use serde_json::json;
2233
2234        let mut json_value = json!({
2235            "extractor": "test_extractor",
2236            "chain": "ethereum",
2237            "block": {
2238                "number": 100,
2239                "hash": "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef",
2240                "parent_hash": "0xabcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890",
2241                "chain": "ethereum",
2242                "ts": "2024-01-01T00:00:00"
2243            },
2244            "finalized_block_height": 99,
2245            "revert": false,
2246            "new_tokens": {},
2247            "account_updates": {},
2248            "state_updates": {},
2249            "new_protocol_components": {},
2250            "deleted_protocol_components": {},
2251            "component_balances": {},
2252            "account_balances": {},
2253            "component_tvl": {},
2254            "dci_update": {
2255                "new_entrypoints": {},
2256                "new_entrypoint_params": {},
2257                "trace_results": {}
2258            }
2259        });
2260
2261        // Add is_partial field only if specified
2262        if let Some(partial_value) = has_partial_value {
2263            json_value["partial_block_index"] = json!(partial_value);
2264        }
2265
2266        let block_changes: BlockChanges =
2267            serde_json::from_value(json_value).expect("Failed to deserialize BlockChanges");
2268
2269        assert_eq!(block_changes.partial_block_index, expected);
2270    }
2271
2272    #[test]
2273    fn test_protocol_components_equality() {
2274        let body1 = ProtocolComponentsRequestBody {
2275            protocol_system: "protocol1".to_string(),
2276            component_ids: Some(vec!["component1".to_string(), "component2".to_string()]),
2277            tvl_gt: Some(1000.0),
2278            chain: Chain::Ethereum,
2279            pagination: PaginationParams::default(),
2280        };
2281
2282        let body2 = ProtocolComponentsRequestBody {
2283            protocol_system: "protocol1".to_string(),
2284            component_ids: Some(vec!["component1".to_string(), "component2".to_string()]),
2285            tvl_gt: Some(1000.0 + 1e-7), // Within the tolerance ±1e-6
2286            chain: Chain::Ethereum,
2287            pagination: PaginationParams::default(),
2288        };
2289
2290        // These should be considered equal due to the tolerance in tvl_gt
2291        assert_eq!(body1, body2);
2292    }
2293
2294    #[test]
2295    fn test_protocol_components_inequality() {
2296        let body1 = ProtocolComponentsRequestBody {
2297            protocol_system: "protocol1".to_string(),
2298            component_ids: Some(vec!["component1".to_string(), "component2".to_string()]),
2299            tvl_gt: Some(1000.0),
2300            chain: Chain::Ethereum,
2301            pagination: PaginationParams::default(),
2302        };
2303
2304        let body2 = ProtocolComponentsRequestBody {
2305            protocol_system: "protocol1".to_string(),
2306            component_ids: Some(vec!["component1".to_string(), "component2".to_string()]),
2307            tvl_gt: Some(1000.0 + 1e-5), // Outside the tolerance ±1e-6
2308            chain: Chain::Ethereum,
2309            pagination: PaginationParams::default(),
2310        };
2311
2312        // These should not be equal due to the difference in tvl_gt
2313        assert_ne!(body1, body2);
2314    }
2315
2316    #[test]
2317    fn test_parse_state_request() {
2318        let json_str = r#"
2319    {
2320        "contractIds": [
2321            "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092"
2322        ],
2323        "protocol_system": "uniswap_v2",
2324        "version": {
2325            "timestamp": "2069-01-01T04:20:00",
2326            "block": {
2327                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2328                "number": 213,
2329                "chain": "ethereum"
2330            }
2331        }
2332    }
2333    "#;
2334
2335        let result: StateRequestBody = serde_json::from_str(json_str).unwrap();
2336
2337        let contract0 = "b4eccE46b8D4e4abFd03C9B806276A6735C9c092"
2338            .parse()
2339            .unwrap();
2340        let block_hash = "24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4"
2341            .parse()
2342            .unwrap();
2343        let block_number = 213;
2344
2345        let expected_timestamp =
2346            NaiveDateTime::parse_from_str("2069-01-01T04:20:00", "%Y-%m-%dT%H:%M:%S").unwrap();
2347
2348        let expected = StateRequestBody {
2349            contract_ids: Some(vec![contract0]),
2350            protocol_system: "uniswap_v2".to_string(),
2351            version: VersionParam {
2352                timestamp: Some(expected_timestamp),
2353                block: Some(BlockParam {
2354                    hash: Some(block_hash),
2355                    chain: Some(Chain::Ethereum),
2356                    number: Some(block_number),
2357                }),
2358            },
2359            chain: Chain::Ethereum,
2360            pagination: PaginationParams::default(),
2361        };
2362
2363        assert_eq!(result, expected);
2364    }
2365
2366    #[test]
2367    fn test_parse_state_request_dual_interface() {
2368        let json_common = r#"
2369    {
2370        "__CONTRACT_IDS__": [
2371            "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092"
2372        ],
2373        "version": {
2374            "timestamp": "2069-01-01T04:20:00",
2375            "block": {
2376                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2377                "number": 213,
2378                "chain": "ethereum"
2379            }
2380        }
2381    }
2382    "#;
2383
2384        let json_str_snake = json_common.replace("\"__CONTRACT_IDS__\"", "\"contract_ids\"");
2385        let json_str_camel = json_common.replace("\"__CONTRACT_IDS__\"", "\"contractIds\"");
2386
2387        let snake: StateRequestBody = serde_json::from_str(&json_str_snake).unwrap();
2388        let camel: StateRequestBody = serde_json::from_str(&json_str_camel).unwrap();
2389
2390        assert_eq!(snake, camel);
2391    }
2392
2393    #[test]
2394    fn test_parse_state_request_unknown_field() {
2395        let body = r#"
2396    {
2397        "contract_ids_with_typo_error": [
2398            {
2399                "address": "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092",
2400                "chain": "ethereum"
2401            }
2402        ],
2403        "version": {
2404            "timestamp": "2069-01-01T04:20:00",
2405            "block": {
2406                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2407                "parentHash": "0x8d75152454e60413efe758cc424bfd339897062d7e658f302765eb7b50971815",
2408                "number": 213,
2409                "chain": "ethereum"
2410            }
2411        }
2412    }
2413    "#;
2414
2415        let decoded = serde_json::from_str::<StateRequestBody>(body);
2416
2417        assert!(decoded.is_err(), "Expected an error due to unknown field");
2418
2419        if let Err(e) = decoded {
2420            assert!(
2421                e.to_string()
2422                    .contains("unknown field `contract_ids_with_typo_error`"),
2423                "Error message does not contain expected unknown field information"
2424            );
2425        }
2426    }
2427
2428    #[test]
2429    fn test_parse_state_request_no_contract_specified() {
2430        let json_str = r#"
2431    {
2432        "protocol_system": "uniswap_v2",
2433        "version": {
2434            "timestamp": "2069-01-01T04:20:00",
2435            "block": {
2436                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2437                "number": 213,
2438                "chain": "ethereum"
2439            }
2440        }
2441    }
2442    "#;
2443
2444        let result: StateRequestBody = serde_json::from_str(json_str).unwrap();
2445
2446        let block_hash = "24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4".into();
2447        let block_number = 213;
2448        let expected_timestamp =
2449            NaiveDateTime::parse_from_str("2069-01-01T04:20:00", "%Y-%m-%dT%H:%M:%S").unwrap();
2450
2451        let expected = StateRequestBody {
2452            contract_ids: None,
2453            protocol_system: "uniswap_v2".to_string(),
2454            version: VersionParam {
2455                timestamp: Some(expected_timestamp),
2456                block: Some(BlockParam {
2457                    hash: Some(block_hash),
2458                    chain: Some(Chain::Ethereum),
2459                    number: Some(block_number),
2460                }),
2461            },
2462            chain: Chain::Ethereum,
2463            pagination: PaginationParams { page: 0, page_size: 100 },
2464        };
2465
2466        assert_eq!(result, expected);
2467    }
2468
2469    #[rstest]
2470    #[case::deprecated_ids(
2471        r#"
2472    {
2473        "protocol_ids": [
2474            {
2475                "id": "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092",
2476                "chain": "ethereum"
2477            }
2478        ],
2479        "protocol_system": "uniswap_v2",
2480        "include_balances": false,
2481        "version": {
2482            "timestamp": "2069-01-01T04:20:00",
2483            "block": {
2484                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2485                "number": 213,
2486                "chain": "ethereum"
2487            }
2488        }
2489    }
2490    "#
2491    )]
2492    #[case(
2493        r#"
2494    {
2495        "protocolIds": [
2496            "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092"
2497        ],
2498        "protocol_system": "uniswap_v2",
2499        "include_balances": false,
2500        "version": {
2501            "timestamp": "2069-01-01T04:20:00",
2502            "block": {
2503                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2504                "number": 213,
2505                "chain": "ethereum"
2506            }
2507        }
2508    }
2509    "#
2510    )]
2511    fn test_parse_protocol_state_request(#[case] json_str: &str) {
2512        let result: ProtocolStateRequestBody = serde_json::from_str(json_str).unwrap();
2513
2514        let block_hash = "24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4"
2515            .parse()
2516            .unwrap();
2517        let block_number = 213;
2518
2519        let expected_timestamp =
2520            NaiveDateTime::parse_from_str("2069-01-01T04:20:00", "%Y-%m-%dT%H:%M:%S").unwrap();
2521
2522        let expected = ProtocolStateRequestBody {
2523            protocol_ids: Some(vec!["0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092".to_string()]),
2524            protocol_system: "uniswap_v2".to_string(),
2525            version: VersionParam {
2526                timestamp: Some(expected_timestamp),
2527                block: Some(BlockParam {
2528                    hash: Some(block_hash),
2529                    chain: Some(Chain::Ethereum),
2530                    number: Some(block_number),
2531                }),
2532            },
2533            chain: Chain::Ethereum,
2534            include_balances: false,
2535            pagination: PaginationParams::default(),
2536        };
2537
2538        assert_eq!(result, expected);
2539    }
2540
2541    #[rstest]
2542    #[case::with_protocol_ids(vec![ProtocolId { id: "id1".to_string(), chain: Chain::Ethereum }, ProtocolId { id: "id2".to_string(), chain: Chain::Ethereum }], vec!["id1".to_string(), "id2".to_string()])]
2543    #[case::with_strings(vec!["id1".to_string(), "id2".to_string()], vec!["id1".to_string(), "id2".to_string()])]
2544    fn test_id_filtered<T>(#[case] input_ids: Vec<T>, #[case] expected_ids: Vec<String>)
2545    where
2546        T: Into<String> + Clone,
2547    {
2548        let request_body = ProtocolStateRequestBody::id_filtered(input_ids);
2549
2550        assert_eq!(request_body.protocol_ids, Some(expected_ids));
2551    }
2552
2553    fn create_models_block_changes() -> crate::models::blockchain::BlockAggregatedChanges {
2554        let base_ts = 1694534400; // Example base timestamp for 2023-09-14T00:00:00
2555
2556        crate::models::blockchain::BlockAggregatedChanges {
2557            extractor: "native_name".to_string(),
2558            block: models::blockchain::Block::new(
2559                3,
2560                models::Chain::Ethereum,
2561                Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000000003").unwrap(),
2562                Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000000002").unwrap(),
2563                chrono::DateTime::from_timestamp(base_ts + 3000, 0).unwrap().naive_utc(),
2564            ),
2565            db_committed_block_height: Some(1),
2566            finalized_block_height: 1,
2567            revert: true,
2568            state_deltas: HashMap::from([
2569                ("pc_1".to_string(), models::protocol::ProtocolComponentStateDelta {
2570                    component_id: "pc_1".to_string(),
2571                    updated_attributes: HashMap::from([
2572                        ("attr_2".to_string(), Bytes::from("0x0000000000000002")),
2573                        ("attr_1".to_string(), Bytes::from("0x00000000000003e8")),
2574                    ]),
2575                    deleted_attributes: HashSet::new(),
2576                }),
2577            ]),
2578            new_protocol_components: HashMap::from([
2579                ("pc_2".to_string(), crate::models::protocol::ProtocolComponent {
2580                    id: "pc_2".to_string(),
2581                    protocol_system: "native_protocol_system".to_string(),
2582                    protocol_type_name: "pt_1".to_string(),
2583                    chain: models::Chain::Ethereum,
2584                    tokens: vec![
2585                        Bytes::from_str("0xdac17f958d2ee523a2206206994597c13d831ec7").unwrap(),
2586                        Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(),
2587                    ],
2588                    contract_addresses: vec![],
2589                    static_attributes: HashMap::new(),
2590                    change: models::ChangeType::Creation,
2591                    creation_tx: Bytes::from_str("0x000000000000000000000000000000000000000000000000000000000000c351").unwrap(),
2592                    created_at: chrono::DateTime::from_timestamp(base_ts + 5000, 0).unwrap().naive_utc(),
2593                }),
2594            ]),
2595            deleted_protocol_components: HashMap::from([
2596                ("pc_3".to_string(), crate::models::protocol::ProtocolComponent {
2597                    id: "pc_3".to_string(),
2598                    protocol_system: "native_protocol_system".to_string(),
2599                    protocol_type_name: "pt_2".to_string(),
2600                    chain: models::Chain::Ethereum,
2601                    tokens: vec![
2602                        Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap(),
2603                        Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(),
2604                    ],
2605                    contract_addresses: vec![],
2606                    static_attributes: HashMap::new(),
2607                    change: models::ChangeType::Deletion,
2608                    creation_tx: Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000009c41").unwrap(),
2609                    created_at: chrono::DateTime::from_timestamp(base_ts + 4000, 0).unwrap().naive_utc(),
2610                }),
2611            ]),
2612            component_balances: HashMap::from([
2613                ("pc_1".to_string(), HashMap::from([
2614                    (Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(), models::protocol::ComponentBalance {
2615                        token: Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(),
2616                        balance: Bytes::from("0x00000001"),
2617                        balance_float: 1.0,
2618                        modify_tx: Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000000000").unwrap(),
2619                        component_id: "pc_1".to_string(),
2620                    }),
2621                    (Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(), models::protocol::ComponentBalance {
2622                        token: Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(),
2623                        balance: Bytes::from("0x000003e8"),
2624                        balance_float: 1000.0,
2625                        modify_tx: Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000007531").unwrap(),
2626                        component_id: "pc_1".to_string(),
2627                    }),
2628                ])),
2629            ]),
2630            account_balances: HashMap::from([
2631                (Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(), HashMap::from([
2632                    (Bytes::from_str("0x7a250d5630b4cf539739df2c5dacb4c659f2488d").unwrap(), models::contract::AccountBalance {
2633                        account: Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(),
2634                        token: Bytes::from_str("0x7a250d5630b4cf539739df2c5dacb4c659f2488d").unwrap(),
2635                        balance: Bytes::from("0x000003e8"),
2636                        modify_tx: Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000007531").unwrap(),
2637                    }),
2638                    ])),
2639            ]),
2640            ..Default::default()
2641        }
2642    }
2643
2644    #[test]
2645    fn test_serialize_deserialize_block_changes() {
2646        // Test that models::BlockAggregatedChanges serialized as json can be deserialized as
2647        // dto::BlockChanges.
2648
2649        // Create a models::BlockAggregatedChanges instance
2650        let block_entity_changes = create_models_block_changes();
2651
2652        // Serialize the struct into JSON
2653        let json_data = serde_json::to_string(&block_entity_changes).expect("Failed to serialize");
2654
2655        // Deserialize the JSON back into a dto::BlockChanges struct
2656        serde_json::from_str::<BlockChanges>(&json_data).expect("parsing failed");
2657    }
2658
2659    #[test]
2660    fn test_parse_block_changes() {
2661        let json_data = r#"
2662        {
2663            "extractor": "vm:ambient",
2664            "chain": "ethereum",
2665            "block": {
2666                "number": 123,
2667                "hash": "0x0000000000000000000000000000000000000000000000000000000000000000",
2668                "parent_hash": "0x0000000000000000000000000000000000000000000000000000000000000000",
2669                "chain": "ethereum",
2670                "ts": "2023-09-14T00:00:00"
2671            },
2672            "finalized_block_height": 0,
2673            "revert": false,
2674            "new_tokens": {},
2675            "account_updates": {
2676                "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2677                    "address": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2678                    "chain": "ethereum",
2679                    "slots": {},
2680                    "balance": "0x01f4",
2681                    "code": "",
2682                    "change": "Update"
2683                }
2684            },
2685            "state_updates": {
2686                "component_1": {
2687                    "component_id": "component_1",
2688                    "updated_attributes": {"attr1": "0x01"},
2689                    "deleted_attributes": ["attr2"]
2690                }
2691            },
2692            "new_protocol_components":
2693                { "protocol_1": {
2694                        "id": "protocol_1",
2695                        "protocol_system": "system_1",
2696                        "protocol_type_name": "type_1",
2697                        "chain": "ethereum",
2698                        "tokens": ["0x01", "0x02"],
2699                        "contract_ids": ["0x01", "0x02"],
2700                        "static_attributes": {"attr1": "0x01f4"},
2701                        "change": "Update",
2702                        "creation_tx": "0x01",
2703                        "created_at": "2023-09-14T00:00:00"
2704                    }
2705                },
2706            "deleted_protocol_components": {},
2707            "component_balances": {
2708                "protocol_1":
2709                    {
2710                        "0x01": {
2711                            "token": "0x01",
2712                            "balance": "0xb77831d23691653a01",
2713                            "balance_float": 3.3844151001790677e21,
2714                            "modify_tx": "0x01",
2715                            "component_id": "protocol_1"
2716                        }
2717                    }
2718            },
2719            "account_balances": {
2720                "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2721                    "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2722                        "account": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2723                        "token": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2724                        "balance": "0x01f4",
2725                        "modify_tx": "0x01"
2726                    }
2727                }
2728            },
2729            "component_tvl": {
2730                "protocol_1": 1000.0
2731            },
2732            "dci_update": {
2733                "new_entrypoints": {
2734                    "component_1": [
2735                        {
2736                            "external_id": "0x01:sig()",
2737                            "target": "0x01",
2738                            "signature": "sig()"
2739                        }
2740                    ]
2741                },
2742                "new_entrypoint_params": {
2743                    "0x01:sig()": [
2744                        [
2745                            {
2746                                "method": "rpctracer",
2747                                "caller": "0x01",
2748                                "calldata": "0x02"
2749                            },
2750                            "component_1"
2751                        ]
2752                    ]
2753                },
2754                "trace_results": {
2755                    "0x01:sig()": {
2756                        "retriggers": [
2757                            ["0x01", {"key": "0x02", "offset": 12}]
2758                        ],
2759                        "accessed_slots": {
2760                            "0x03": ["0x03", "0x04"]
2761                        }
2762                    }
2763                }
2764            }
2765        }
2766        "#;
2767
2768        serde_json::from_str::<BlockChanges>(json_data).expect("parsing failed");
2769    }
2770
2771    #[test]
2772    fn test_parse_websocket_message() {
2773        let json_data = r#"
2774        {
2775            "subscription_id": "5d23bfbe-89ad-4ea3-8672-dc9e973ac9dc",
2776            "deltas": {
2777                "type": "BlockChanges",
2778                "extractor": "uniswap_v2",
2779                "chain": "ethereum",
2780                "block": {
2781                    "number": 19291517,
2782                    "hash": "0xbc3ea4896c0be8da6229387a8571b72818aa258daf4fab46471003ad74c4ee83",
2783                    "parent_hash": "0x89ca5b8d593574cf6c886f41ef8208bf6bdc1a90ef36046cb8c84bc880b9af8f",
2784                    "chain": "ethereum",
2785                    "ts": "2024-02-23T16:35:35"
2786                },
2787                "finalized_block_height": 0,
2788                "revert": false,
2789                "new_tokens": {},
2790                "account_updates": {
2791                    "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2792                        "address": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2793                        "chain": "ethereum",
2794                        "slots": {},
2795                        "balance": "0x01f4",
2796                        "code": "",
2797                        "change": "Update"
2798                    }
2799                },
2800                "state_updates": {
2801                    "0xde6faedbcae38eec6d33ad61473a04a6dd7f6e28": {
2802                        "component_id": "0xde6faedbcae38eec6d33ad61473a04a6dd7f6e28",
2803                        "updated_attributes": {
2804                            "reserve0": "0x87f7b5973a7f28a8b32404",
2805                            "reserve1": "0x09e9564b11"
2806                        },
2807                        "deleted_attributes": []
2808                    },
2809                    "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d": {
2810                        "component_id": "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d",
2811                        "updated_attributes": {
2812                            "reserve1": "0x44d9a8fd662c2f4d03",
2813                            "reserve0": "0x500b1261f811d5bf423e"
2814                        },
2815                        "deleted_attributes": []
2816                    }
2817                },
2818                "new_protocol_components": {},
2819                "deleted_protocol_components": {},
2820                "component_balances": {
2821                    "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d": {
2822                        "0x9012744b7a564623b6c3e40b144fc196bdedf1a9": {
2823                            "token": "0x9012744b7a564623b6c3e40b144fc196bdedf1a9",
2824                            "balance": "0x500b1261f811d5bf423e",
2825                            "balance_float": 3.779935574269033E23,
2826                            "modify_tx": "0xe46c4db085fb6c6f3408a65524555797adb264e1d5cf3b66ad154598f85ac4bf",
2827                            "component_id": "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d"
2828                        },
2829                        "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2": {
2830                            "token": "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2",
2831                            "balance": "0x44d9a8fd662c2f4d03",
2832                            "balance_float": 1.270062661329837E21,
2833                            "modify_tx": "0xe46c4db085fb6c6f3408a65524555797adb264e1d5cf3b66ad154598f85ac4bf",
2834                            "component_id": "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d"
2835                        }
2836                    }
2837                },
2838                "account_balances": {
2839                    "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2840                        "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2841                            "account": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2842                            "token": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2843                            "balance": "0x01f4",
2844                            "modify_tx": "0x01"
2845                        }
2846                    }
2847                },
2848                "component_tvl": {},
2849                "dci_update": {
2850                    "new_entrypoints": {
2851                        "0xde6faedbcae38eec6d33ad61473a04a6dd7f6e28": [
2852                            {
2853                                "external_id": "0x01:sig()",
2854                                "target": "0x01",
2855                                "signature": "sig()"
2856                            }
2857                        ]
2858                    },
2859                    "new_entrypoint_params": {
2860                        "0x01:sig()": [
2861                            [
2862                                {
2863                                    "method": "rpctracer",
2864                                    "caller": "0x01",
2865                                    "calldata": "0x02"
2866                                },
2867                                "0xde6faedbcae38eec6d33ad61473a04a6dd7f6e28"
2868                            ]
2869                        ]
2870                    },
2871                    "trace_results": {
2872                        "0x01:sig()": {
2873                            "retriggers": [
2874                                ["0x01", {"key": "0x02", "offset": 12}]
2875                            ],
2876                            "accessed_slots": {
2877                                "0x03": ["0x03", "0x04"]
2878                            }
2879                        }
2880                    }
2881                }
2882            }
2883        }
2884        "#;
2885        serde_json::from_str::<WebSocketMessage>(json_data).expect("parsing failed");
2886    }
2887
2888    #[test]
2889    fn test_protocol_state_delta_merge_update_delete() {
2890        // Initialize ProtocolStateDelta instances
2891        let mut delta1 = ProtocolStateDelta {
2892            component_id: "Component1".to_string(),
2893            updated_attributes: HashMap::from([(
2894                "Attribute1".to_string(),
2895                Bytes::from("0xbadbabe420"),
2896            )]),
2897            deleted_attributes: HashSet::new(),
2898        };
2899        let delta2 = ProtocolStateDelta {
2900            component_id: "Component1".to_string(),
2901            updated_attributes: HashMap::from([(
2902                "Attribute2".to_string(),
2903                Bytes::from("0x0badbabe"),
2904            )]),
2905            deleted_attributes: HashSet::from(["Attribute1".to_string()]),
2906        };
2907        let exp = ProtocolStateDelta {
2908            component_id: "Component1".to_string(),
2909            updated_attributes: HashMap::from([(
2910                "Attribute2".to_string(),
2911                Bytes::from("0x0badbabe"),
2912            )]),
2913            deleted_attributes: HashSet::from(["Attribute1".to_string()]),
2914        };
2915
2916        delta1.merge(&delta2);
2917
2918        assert_eq!(delta1, exp);
2919    }
2920
2921    #[test]
2922    fn test_protocol_state_delta_merge_delete_update() {
2923        // Initialize ProtocolStateDelta instances
2924        let mut delta1 = ProtocolStateDelta {
2925            component_id: "Component1".to_string(),
2926            updated_attributes: HashMap::new(),
2927            deleted_attributes: HashSet::from(["Attribute1".to_string()]),
2928        };
2929        let delta2 = ProtocolStateDelta {
2930            component_id: "Component1".to_string(),
2931            updated_attributes: HashMap::from([(
2932                "Attribute1".to_string(),
2933                Bytes::from("0x0badbabe"),
2934            )]),
2935            deleted_attributes: HashSet::new(),
2936        };
2937        let exp = ProtocolStateDelta {
2938            component_id: "Component1".to_string(),
2939            updated_attributes: HashMap::from([(
2940                "Attribute1".to_string(),
2941                Bytes::from("0x0badbabe"),
2942            )]),
2943            deleted_attributes: HashSet::new(),
2944        };
2945
2946        delta1.merge(&delta2);
2947
2948        assert_eq!(delta1, exp);
2949    }
2950
2951    #[test]
2952    fn test_account_update_merge() {
2953        // Initialize AccountUpdate instances with same address and valid hex strings for Bytes
2954        let mut account1 = AccountUpdate::new(
2955            Bytes::from(b"0x1234"),
2956            Chain::Ethereum,
2957            HashMap::from([(Bytes::from("0xaabb"), Bytes::from("0xccdd"))]),
2958            Some(Bytes::from("0x1000")),
2959            Some(Bytes::from("0xdeadbeaf")),
2960            ChangeType::Creation,
2961        );
2962
2963        let account2 = AccountUpdate::new(
2964            Bytes::from(b"0x1234"), // Same id as account1
2965            Chain::Ethereum,
2966            HashMap::from([(Bytes::from("0xeeff"), Bytes::from("0x11223344"))]),
2967            Some(Bytes::from("0x2000")),
2968            Some(Bytes::from("0xcafebabe")),
2969            ChangeType::Update,
2970        );
2971
2972        // Merge account2 into account1
2973        account1.merge(&account2);
2974
2975        // Define the expected state after merge
2976        let expected = AccountUpdate::new(
2977            Bytes::from(b"0x1234"), // Same id as before the merge
2978            Chain::Ethereum,
2979            HashMap::from([
2980                (Bytes::from("0xaabb"), Bytes::from("0xccdd")), // Original slot from account1
2981                (Bytes::from("0xeeff"), Bytes::from("0x11223344")), // New slot from account2
2982            ]),
2983            Some(Bytes::from("0x2000")),     // Updated balance
2984            Some(Bytes::from("0xcafebabe")), // Updated code
2985            ChangeType::Creation,            // Updated change type
2986        );
2987
2988        // Assert the new account1 equals to the expected state
2989        assert_eq!(account1, expected);
2990    }
2991
2992    #[test]
2993    fn test_block_account_changes_merge() {
2994        // Prepare account updates
2995        let old_account_updates: HashMap<Bytes, AccountUpdate> = [(
2996            Bytes::from("0x0011"),
2997            AccountUpdate {
2998                address: Bytes::from("0x00"),
2999                chain: Chain::Ethereum,
3000                slots: HashMap::from([(Bytes::from("0x0022"), Bytes::from("0x0033"))]),
3001                balance: Some(Bytes::from("0x01")),
3002                code: Some(Bytes::from("0x02")),
3003                change: ChangeType::Creation,
3004            },
3005        )]
3006        .into_iter()
3007        .collect();
3008        let new_account_updates: HashMap<Bytes, AccountUpdate> = [(
3009            Bytes::from("0x0011"),
3010            AccountUpdate {
3011                address: Bytes::from("0x00"),
3012                chain: Chain::Ethereum,
3013                slots: HashMap::from([(Bytes::from("0x0044"), Bytes::from("0x0055"))]),
3014                balance: Some(Bytes::from("0x03")),
3015                code: Some(Bytes::from("0x04")),
3016                change: ChangeType::Update,
3017            },
3018        )]
3019        .into_iter()
3020        .collect();
3021        // Create initial and new BlockAccountChanges instances
3022        let block_account_changes_initial = BlockChanges {
3023            extractor: "extractor1".to_string(),
3024            revert: false,
3025            account_updates: old_account_updates,
3026            ..Default::default()
3027        };
3028
3029        let block_account_changes_new = BlockChanges {
3030            extractor: "extractor2".to_string(),
3031            revert: true,
3032            account_updates: new_account_updates,
3033            ..Default::default()
3034        };
3035
3036        // Merge the new BlockChanges into the initial one
3037        let res = block_account_changes_initial.merge(block_account_changes_new);
3038
3039        // Create the expected result of the merge operation
3040        let expected_account_updates: HashMap<Bytes, AccountUpdate> = [(
3041            Bytes::from("0x0011"),
3042            AccountUpdate {
3043                address: Bytes::from("0x00"),
3044                chain: Chain::Ethereum,
3045                slots: HashMap::from([
3046                    (Bytes::from("0x0044"), Bytes::from("0x0055")),
3047                    (Bytes::from("0x0022"), Bytes::from("0x0033")),
3048                ]),
3049                balance: Some(Bytes::from("0x03")),
3050                code: Some(Bytes::from("0x04")),
3051                change: ChangeType::Creation,
3052            },
3053        )]
3054        .into_iter()
3055        .collect();
3056        let block_account_changes_expected = BlockChanges {
3057            extractor: "extractor1".to_string(),
3058            revert: true,
3059            account_updates: expected_account_updates,
3060            ..Default::default()
3061        };
3062        assert_eq!(res, block_account_changes_expected);
3063    }
3064
3065    #[test]
3066    fn test_block_entity_changes_merge() {
3067        // Initialize two BlockChanges instances with different details
3068        let block_entity_changes_result1 = BlockChanges {
3069            extractor: String::from("extractor1"),
3070            revert: false,
3071            state_updates: hashmap! { "state1".to_string() => ProtocolStateDelta::default() },
3072            new_protocol_components: hashmap! { "component1".to_string() => ProtocolComponent::default() },
3073            deleted_protocol_components: HashMap::new(),
3074            component_balances: hashmap! {
3075                "component1".to_string() => TokenBalances(hashmap! {
3076                    Bytes::from("0x01") => ComponentBalance {
3077                            token: Bytes::from("0x01"),
3078                            balance: Bytes::from("0x01"),
3079                            balance_float: 1.0,
3080                            modify_tx: Bytes::from("0x00"),
3081                            component_id: "component1".to_string()
3082                        },
3083                    Bytes::from("0x02") => ComponentBalance {
3084                        token: Bytes::from("0x02"),
3085                        balance: Bytes::from("0x02"),
3086                        balance_float: 2.0,
3087                        modify_tx: Bytes::from("0x00"),
3088                        component_id: "component1".to_string()
3089                    },
3090                })
3091
3092            },
3093            component_tvl: hashmap! { "tvl1".to_string() => 1000.0 },
3094            ..Default::default()
3095        };
3096        let block_entity_changes_result2 = BlockChanges {
3097            extractor: String::from("extractor2"),
3098            revert: true,
3099            state_updates: hashmap! { "state2".to_string() => ProtocolStateDelta::default() },
3100            new_protocol_components: hashmap! { "component2".to_string() => ProtocolComponent::default() },
3101            deleted_protocol_components: hashmap! { "component3".to_string() => ProtocolComponent::default() },
3102            component_balances: hashmap! {
3103                "component1".to_string() => TokenBalances::default(),
3104                "component2".to_string() => TokenBalances::default()
3105            },
3106            component_tvl: hashmap! { "tvl2".to_string() => 2000.0 },
3107            ..Default::default()
3108        };
3109
3110        let res = block_entity_changes_result1.merge(block_entity_changes_result2);
3111
3112        let expected_block_entity_changes_result = BlockChanges {
3113            extractor: String::from("extractor1"),
3114            revert: true,
3115            state_updates: hashmap! {
3116                "state1".to_string() => ProtocolStateDelta::default(),
3117                "state2".to_string() => ProtocolStateDelta::default(),
3118            },
3119            new_protocol_components: hashmap! {
3120                "component1".to_string() => ProtocolComponent::default(),
3121                "component2".to_string() => ProtocolComponent::default(),
3122            },
3123            deleted_protocol_components: hashmap! {
3124                "component3".to_string() => ProtocolComponent::default(),
3125            },
3126            component_balances: hashmap! {
3127                "component1".to_string() => TokenBalances(hashmap! {
3128                    Bytes::from("0x01") => ComponentBalance {
3129                            token: Bytes::from("0x01"),
3130                            balance: Bytes::from("0x01"),
3131                            balance_float: 1.0,
3132                            modify_tx: Bytes::from("0x00"),
3133                            component_id: "component1".to_string()
3134                        },
3135                    Bytes::from("0x02") => ComponentBalance {
3136                        token: Bytes::from("0x02"),
3137                        balance: Bytes::from("0x02"),
3138                        balance_float: 2.0,
3139                        modify_tx: Bytes::from("0x00"),
3140                        component_id: "component1".to_string()
3141                        },
3142                    }),
3143                "component2".to_string() => TokenBalances::default(),
3144            },
3145            component_tvl: hashmap! {
3146                "tvl1".to_string() => 1000.0,
3147                "tvl2".to_string() => 2000.0
3148            },
3149            ..Default::default()
3150        };
3151
3152        assert_eq!(res, expected_block_entity_changes_result);
3153    }
3154
3155    #[test]
3156    fn test_websocket_error_serialization() {
3157        let extractor_id = ExtractorIdentity::new(Chain::Ethereum, "test_extractor");
3158        let subscription_id = Uuid::new_v4();
3159
3160        // Test ExtractorNotFound serialization
3161        let error = WebsocketError::ExtractorNotFound(extractor_id.clone());
3162        let json = serde_json::to_string(&error).unwrap();
3163        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3164        assert_eq!(error, deserialized);
3165
3166        // Test SubscriptionNotFound serialization
3167        let error = WebsocketError::SubscriptionNotFound(subscription_id);
3168        let json = serde_json::to_string(&error).unwrap();
3169        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3170        assert_eq!(error, deserialized);
3171
3172        // Test ParseError serialization
3173        let error = WebsocketError::ParseError("{asd".to_string(), "invalid json".to_string());
3174        let json = serde_json::to_string(&error).unwrap();
3175        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3176        assert_eq!(error, deserialized);
3177
3178        // Test SubscribeError serialization
3179        let error = WebsocketError::SubscribeError(extractor_id.clone());
3180        let json = serde_json::to_string(&error).unwrap();
3181        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3182        assert_eq!(error, deserialized);
3183
3184        // Test CompressionError serialization
3185        let error =
3186            WebsocketError::CompressionError(subscription_id, "Compression failed".to_string());
3187        let json = serde_json::to_string(&error).unwrap();
3188        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3189        assert_eq!(error, deserialized);
3190    }
3191
3192    #[test]
3193    fn test_websocket_message_with_error_response() {
3194        let error =
3195            WebsocketError::ParseError("}asdfas".to_string(), "malformed request".to_string());
3196        let response = Response::Error(error.clone());
3197        let message = WebSocketMessage::Response(response);
3198
3199        let json = serde_json::to_string(&message).unwrap();
3200        let deserialized: WebSocketMessage = serde_json::from_str(&json).unwrap();
3201
3202        if let WebSocketMessage::Response(Response::Error(deserialized_error)) = deserialized {
3203            assert_eq!(error, deserialized_error);
3204        } else {
3205            panic!("Expected WebSocketMessage::Response(Response::Error)");
3206        }
3207    }
3208
3209    #[test]
3210    fn test_websocket_error_conversion_from_models() {
3211        use crate::models::error::WebsocketError as ModelsError;
3212
3213        let extractor_id =
3214            crate::models::ExtractorIdentity::new(crate::models::Chain::Ethereum, "test");
3215        let subscription_id = Uuid::new_v4();
3216
3217        // Test ExtractorNotFound conversion
3218        let models_error = ModelsError::ExtractorNotFound(extractor_id.clone());
3219        let dto_error: WebsocketError = models_error.into();
3220        assert_eq!(dto_error, WebsocketError::ExtractorNotFound(extractor_id.clone().into()));
3221
3222        // Test SubscriptionNotFound conversion
3223        let models_error = ModelsError::SubscriptionNotFound(subscription_id);
3224        let dto_error: WebsocketError = models_error.into();
3225        assert_eq!(dto_error, WebsocketError::SubscriptionNotFound(subscription_id));
3226
3227        // Test ParseError conversion - create a real JSON parse error
3228        let json_result: Result<serde_json::Value, _> = serde_json::from_str("{invalid json");
3229        let json_error = json_result.unwrap_err();
3230        let models_error = ModelsError::ParseError("{invalid json".to_string(), json_error);
3231        let dto_error: WebsocketError = models_error.into();
3232        if let WebsocketError::ParseError(msg, error) = dto_error {
3233            // Just check that we have a non-empty error message
3234            assert!(!error.is_empty(), "Error message should not be empty, got: '{}'", msg);
3235        } else {
3236            panic!("Expected ParseError variant");
3237        }
3238
3239        // Test SubscribeError conversion
3240        let models_error = ModelsError::SubscribeError(extractor_id.clone());
3241        let dto_error: WebsocketError = models_error.into();
3242        assert_eq!(dto_error, WebsocketError::SubscribeError(extractor_id.into()));
3243
3244        // Test CompressionError conversion
3245        let io_error = std::io::Error::other("Compression failed");
3246        let models_error = ModelsError::CompressionError(subscription_id, io_error);
3247        let dto_error: WebsocketError = models_error.into();
3248        if let WebsocketError::CompressionError(sub_id, msg) = &dto_error {
3249            assert_eq!(*sub_id, subscription_id);
3250            assert!(msg.contains("Compression failed"));
3251        } else {
3252            panic!("Expected CompressionError variant");
3253        }
3254    }
3255}
3256
3257#[cfg(test)]
3258mod memory_size_tests {
3259    use std::collections::HashMap;
3260
3261    use super::*;
3262
3263    #[test]
3264    fn test_state_request_response_memory_size_empty() {
3265        let response = StateRequestResponse {
3266            accounts: vec![],
3267            pagination: PaginationResponse::new(1, 10, 0),
3268        };
3269
3270        let size = response.deep_size_of();
3271
3272        // Should at least include base struct sizes
3273        assert!(size >= 48, "Empty response should have minimum size of 48 bytes, got {}", size);
3274        assert!(size < 200, "Empty response should not be too large, got {}", size);
3275    }
3276
3277    #[test]
3278    fn test_state_request_response_memory_size_scales_with_slots() {
3279        let create_response_with_slots = |slot_count: usize| {
3280            let mut slots = HashMap::new();
3281            for i in 0..slot_count {
3282                let key = vec![i as u8; 32]; // 32-byte key
3283                let value = vec![(i + 100) as u8; 32]; // 32-byte value
3284                slots.insert(key.into(), value.into());
3285            }
3286
3287            let account = ResponseAccount::new(
3288                Chain::Ethereum,
3289                vec![1; 20].into(),
3290                "Pool".to_string(),
3291                slots,
3292                vec![1; 32].into(),
3293                HashMap::new(),
3294                vec![].into(), // empty code
3295                vec![1; 32].into(),
3296                vec![1; 32].into(),
3297                vec![1; 32].into(),
3298                None,
3299            );
3300
3301            StateRequestResponse {
3302                accounts: vec![account],
3303                pagination: PaginationResponse::new(1, 10, 1),
3304            }
3305        };
3306
3307        let small_response = create_response_with_slots(10);
3308        let large_response = create_response_with_slots(100);
3309
3310        let small_size = small_response.deep_size_of();
3311        let large_size = large_response.deep_size_of();
3312
3313        // Large response should be significantly bigger
3314        assert!(
3315            large_size > small_size * 5,
3316            "Large response ({} bytes) should be much larger than small response ({} bytes)",
3317            large_size,
3318            small_size
3319        );
3320
3321        // Each slot should contribute at least 64 bytes (32 + 32 + overhead)
3322        let size_diff = large_size - small_size;
3323        let expected_min_diff = 90 * 64; // 90 additional slots * 64 bytes each
3324        assert!(
3325            size_diff > expected_min_diff,
3326            "Size difference ({} bytes) should reflect the additional slot data",
3327            size_diff
3328        );
3329    }
3330}
3331
3332#[cfg(test)]
3333mod pagination_limits_tests {
3334    use super::*;
3335
3336    // Test struct for pagination limits
3337    #[derive(Clone, Debug)]
3338    struct TestRequestBody {
3339        pagination: PaginationParams,
3340    }
3341
3342    // Implement pagination limits for test struct
3343    impl_pagination_limits!(TestRequestBody, compressed = 500, uncompressed = 50);
3344
3345    #[test]
3346    fn test_effective_max_page_size() {
3347        // Test effective max with compression enabled
3348        let max_size = TestRequestBody::effective_max_page_size(true);
3349        assert_eq!(max_size, 500, "Should return compressed limit when compression is enabled");
3350
3351        // Test effective max with compression disabled
3352        let max_size = TestRequestBody::effective_max_page_size(false);
3353        assert_eq!(max_size, 50, "Should return uncompressed limit when compression is disabled");
3354    }
3355}