drt_sc_snippets/
account_tool.rs

1use drt_chain_scenario_format::interpret_trait::IntoRaw;
2use drt_sc_scenario::{
3    imports::Bech32Address,
4    scenario_model::{Account, BytesKey, BytesValue, Scenario, SetStateStep, Step},
5};
6use drt_sdk::{
7    data::{address::Address, dcdt::DcdtBalance},
8    gateway::GatewayProxy,
9};
10use std::collections::{BTreeMap, HashMap};
11
12/// Called directly from CLI, from `sc-meta`.
13///
14/// Retrieves an account data via the API,
15/// then formats it as a scenario set state step.
16pub async fn print_account_as_scenario_set_state(
17    api_string: String,
18    address_bech32_string: String,
19) {
20    let api = GatewayProxy::new(api_string);
21    let address = Bech32Address::from_bech32_string(address_bech32_string);
22    let set_state = retrieve_account_as_scenario_set_state(&api, &address).await;
23    let scenario = build_scenario(set_state);
24    println!("{}", scenario.into_raw().to_json_string());
25}
26
27fn build_scenario(set_state: SetStateStep) -> Scenario {
28    Scenario {
29        name: None,
30        comment: None,
31        check_gas: None,
32        steps: vec![Step::SetState(set_state)],
33    }
34}
35
36pub async fn retrieve_account_as_scenario_set_state(
37    api: &GatewayProxy,
38    address: &Bech32Address,
39) -> SetStateStep {
40    let sdk_address = Address::from_bech32_string(address.to_bech32_str()).unwrap();
41    let sdk_account = api.get_account(&sdk_address).await.unwrap();
42
43    let account_dcdt = api
44        .get_account_dcdt_tokens(&sdk_address)
45        .await
46        .unwrap_or_else(|err| {
47            panic!("failed to retrieve DCDT tokens for address {address}: {err}")
48        });
49    let account_dcdt_roles = api
50        .get_account_dcdt_roles(&sdk_address)
51        .await
52        .unwrap_or_else(|err| panic!("failed to retrieve DCDT roles for address {address}: {err}"));
53    let account_storage = api
54        .get_account_storage_keys(&sdk_address)
55        .await
56        .unwrap_or_else(|err| panic!("failed to retrieve storage for address {address}: {err}"));
57
58    let account_state = set_account(
59        sdk_account,
60        account_storage,
61        account_dcdt,
62        account_dcdt_roles,
63    );
64
65    let set_state_step = SetStateStep::new();
66    set_state_step.put_account(address, account_state)
67}
68
69fn set_account(
70    account: drt_sdk::data::account::Account,
71    account_storage: HashMap<String, String>,
72    account_dcdt: HashMap<String, DcdtBalance>,
73    account_dcdt_roles: HashMap<String, Vec<String>>,
74) -> Account {
75    let mut account_state = Account::new()
76        .nonce(account.nonce)
77        .balance(account.balance.as_str())
78        .code(account.code);
79    account_state.username = Some(format!("str:{}", account.username.as_str()).into());
80    account_state.storage = convert_storage(account_storage);
81
82    for (_, dcdt_balance) in account_dcdt.iter() {
83        let token_id_expr = format!("str:{}", dcdt_balance.token_identifier);
84        account_state =
85            account_state.dcdt_balance(token_id_expr.as_str(), dcdt_balance.balance.as_str());
86    }
87
88    for (token_id, dcdt_roles) in account_dcdt_roles {
89        let token_id_expr = format!("str:{token_id}");
90        account_state = account_state.dcdt_roles(token_id_expr.as_str(), dcdt_roles);
91    }
92
93    account_state
94}
95
96fn convert_storage(account_storage: HashMap<String, String>) -> BTreeMap<BytesKey, BytesValue> {
97    account_storage
98        .into_iter()
99        .filter(|(k, _)| !k.starts_with("454c524f4e44"))
100        .map(|(k, v)| (BytesKey::from(k.as_str()), BytesValue::from(v)))
101        .collect()
102}