nautilus_serialization/arrow/
account_state.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use nautilus_model::events::AccountState;
17
18use super::json::{JsonFieldSpec, impl_json_arrow};
19
20const ACCOUNT_STATE_FIELDS: &[JsonFieldSpec] = &[
21    JsonFieldSpec::utf8("account_id", false),
22    JsonFieldSpec::utf8("account_type", false),
23    JsonFieldSpec::utf8("base_currency", true),
24    JsonFieldSpec::utf8_json("balances", false),
25    JsonFieldSpec::utf8_json("margins", false),
26    JsonFieldSpec::boolean("is_reported", false),
27    JsonFieldSpec::utf8("event_id", false),
28    JsonFieldSpec::u64("ts_event", false),
29    JsonFieldSpec::u64("ts_init", false),
30    JsonFieldSpec::utf8_json("info", true),
31];
32
33impl_json_arrow!(typed AccountState, "AccountState", ACCOUNT_STATE_FIELDS, &["info"]);
34
35#[cfg(test)]
36mod tests {
37    use nautilus_core::Params;
38    use nautilus_model::events::account::stubs::cash_account_state;
39    use rstest::rstest;
40    use serde_json::json;
41
42    use super::*;
43    use crate::arrow::{DecodeTypedFromRecordBatch, EncodeToRecordBatch, json::encode_batch};
44
45    #[rstest]
46    fn test_account_state_round_trip(cash_account_state: AccountState) {
47        let mut info = Params::new();
48        info.insert(
49            "total_wallet_balance".to_string(),
50            json!("1525000.00000001"),
51        );
52        info.insert("can_trade".to_string(), json!(true));
53        let state = cash_account_state.with_info(Some(info));
54        let metadata = state.metadata();
55        let batch = AccountState::encode_batch(&metadata, std::slice::from_ref(&state)).unwrap();
56        let decoded = AccountState::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
57
58        assert_eq!(decoded.len(), 1);
59        assert_eq!(decoded[0].account_id, state.account_id);
60        assert_eq!(decoded[0].balances, state.balances);
61        assert_eq!(decoded[0].margins, state.margins);
62        assert_eq!(decoded[0].base_currency, state.base_currency);
63        assert_eq!(decoded[0].info, state.info);
64    }
65
66    #[rstest]
67    fn test_account_state_decodes_legacy_batch_without_info(cash_account_state: AccountState) {
68        let metadata = cash_account_state.metadata();
69        let legacy_fields = &ACCOUNT_STATE_FIELDS[..ACCOUNT_STATE_FIELDS.len() - 1];
70        let batch = encode_batch(
71            "AccountState",
72            &metadata,
73            std::slice::from_ref(&cash_account_state),
74            legacy_fields,
75        )
76        .unwrap();
77        let decoded = AccountState::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
78
79        assert_eq!(decoded.len(), 1);
80        assert!(decoded[0].info.is_none());
81    }
82}
nautilus_serialization/arrow/account_state.rs

nautilus_serialization/arrow/
account_state.rs
