barter_execution/exchange/mock/
account.rs

1use crate::{
2    UnindexedAccountSnapshot,
3    balance::AssetBalance,
4    order::{
5        Order,
6        id::ClientOrderId,
7        state::{ActiveOrderState, Cancelled, InactiveOrderState, Open, OrderState},
8    },
9    trade::Trade,
10};
11use barter_instrument::{
12    asset::{QuoteAsset, name::AssetNameExchange},
13    exchange::ExchangeId,
14    instrument::name::InstrumentNameExchange,
15};
16use chrono::{DateTime, Utc};
17use derive_more::Constructor;
18use fnv::FnvHashMap;
19
20#[derive(Debug, Constructor)]
21pub struct AccountState {
22    balances: FnvHashMap<AssetNameExchange, AssetBalance<AssetNameExchange>>,
23    orders_open: FnvHashMap<ClientOrderId, Order<ExchangeId, InstrumentNameExchange, Open>>,
24    orders_cancelled:
25        FnvHashMap<ClientOrderId, Order<ExchangeId, InstrumentNameExchange, Cancelled>>,
26    trades: Vec<Trade<QuoteAsset, InstrumentNameExchange>>,
27}
28
29impl AccountState {
30    pub fn update_time_exchange(&mut self, time_exchange: DateTime<Utc>) {
31        for balance in self.balances.values_mut() {
32            balance.time_exchange = time_exchange;
33        }
34
35        for order in self.orders_open.values_mut() {
36            order.state.time_exchange = time_exchange;
37        }
38    }
39
40    pub fn balances(&self) -> impl Iterator<Item = &AssetBalance<AssetNameExchange>> + '_ {
41        self.balances.values()
42    }
43
44    pub fn orders_open(
45        &self,
46    ) -> impl Iterator<Item = &Order<ExchangeId, InstrumentNameExchange, Open>> + '_ {
47        self.orders_open.values()
48    }
49
50    pub fn orders_cancelled(
51        &self,
52    ) -> impl Iterator<Item = &Order<ExchangeId, InstrumentNameExchange, Cancelled>> + '_ {
53        self.orders_cancelled.values()
54    }
55
56    pub fn trades(
57        &self,
58        time_since: DateTime<Utc>,
59    ) -> impl Iterator<Item = &Trade<QuoteAsset, InstrumentNameExchange>> + '_ {
60        self.trades
61            .iter()
62            .filter(move |trade| trade.time_exchange >= time_since)
63    }
64
65    pub fn balance_mut(
66        &mut self,
67        asset: &AssetNameExchange,
68    ) -> Option<&mut AssetBalance<AssetNameExchange>> {
69        self.balances.get_mut(asset)
70    }
71
72    pub fn ack_trade(&mut self, trade: Trade<QuoteAsset, InstrumentNameExchange>) {
73        self.trades.push(trade);
74    }
75}
76
77impl From<UnindexedAccountSnapshot> for AccountState {
78    fn from(value: UnindexedAccountSnapshot) -> Self {
79        let UnindexedAccountSnapshot {
80            exchange: _,
81            balances,
82            instruments,
83        } = value;
84
85        let balances = balances
86            .into_iter()
87            .map(|asset_balance| (asset_balance.asset.clone(), asset_balance))
88            .collect();
89
90        let (orders_open, orders_cancelled) = instruments.into_iter().fold(
91            (FnvHashMap::default(), FnvHashMap::default()),
92            |(mut orders_open, mut orders_cancelled), snapshot| {
93                for order in snapshot.orders {
94                    match order.state {
95                        OrderState::Active(ActiveOrderState::Open(open)) => {
96                            orders_open.insert(
97                                order.key.cid.clone(),
98                                Order {
99                                    key: order.key,
100                                    side: order.side,
101                                    price: order.price,
102                                    quantity: order.quantity,
103                                    kind: order.kind,
104                                    time_in_force: order.time_in_force,
105                                    state: open,
106                                },
107                            );
108                        }
109                        OrderState::Inactive(InactiveOrderState::Cancelled(cancelled)) => {
110                            orders_cancelled.insert(
111                                order.key.cid.clone(),
112                                Order {
113                                    key: order.key,
114                                    side: order.side,
115                                    price: order.price,
116                                    quantity: order.quantity,
117                                    kind: order.kind,
118                                    time_in_force: order.time_in_force,
119                                    state: cancelled,
120                                },
121                            );
122                        }
123                        _ => {}
124                    }
125                }
126
127                (orders_open, orders_cancelled)
128            },
129        );
130
131        Self {
132            balances,
133            orders_open,
134            orders_cancelled,
135            trades: vec![],
136        }
137    }
138}