transaction_fuzzer/account_universe/
account.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4// Copyright (c) The Diem Core Contributors
5// SPDX-License-Identifier: Apache-2.0
6
7use std::sync::Arc;
8
9use proptest::prelude::*;
10use sui_types::{
11    base_types::{ObjectID, SuiAddress},
12    crypto::{AccountKeyPair, get_key_pair},
13    object::Object,
14};
15
16use crate::executor::Executor;
17
18pub const INITIAL_BALANCE: u64 = 100_000_000_000_000;
19pub const PUBLISH_BUDGET: u64 = 1_000_000_000;
20pub const NUM_GAS_OBJECTS: usize = 1;
21
22#[derive(Debug)]
23pub struct Account {
24    pub address: SuiAddress,
25    pub key: AccountKeyPair,
26}
27
28// `Arc` account since the key pair is non-copyable
29#[derive(Debug, Clone)]
30pub struct AccountData {
31    pub account: Arc<Account>,
32    pub coins: Vec<Object>,
33    pub initial_balances: Vec<u64>,
34    pub balance_creation_amt: u64,
35}
36
37#[derive(Clone, Debug)]
38pub struct AccountCurrent {
39    pub initial_data: AccountData,
40    pub current_balances: Vec<u64>,
41    pub current_coins: Vec<Object>,
42    // Non-coin objects
43    pub current_objects: Vec<ObjectID>,
44}
45
46impl Account {
47    pub fn new_random() -> Self {
48        let (address, key) = get_key_pair();
49        Self { address, key }
50    }
51}
52
53impl AccountData {
54    pub fn new_random() -> Self {
55        let account = Account::new_random();
56        Self::new_with_account_and_balance(Arc::new(account), INITIAL_BALANCE)
57    }
58
59    pub fn new_with_account_and_balance(account: Arc<Account>, initial_balance: u64) -> Self {
60        let coins = (0..NUM_GAS_OBJECTS)
61            .map(|_| {
62                let gas_object_id = ObjectID::random();
63                Object::with_id_owner_gas_for_testing(
64                    gas_object_id,
65                    account.address,
66                    initial_balance,
67                )
68            })
69            .collect();
70        let initial_balances = (0..NUM_GAS_OBJECTS).map(|_| initial_balance).collect();
71        Self {
72            account,
73            coins,
74            initial_balances,
75            balance_creation_amt: initial_balance,
76        }
77    }
78}
79
80impl AccountCurrent {
81    pub fn new(account: AccountData) -> Self {
82        Self {
83            current_balances: account.initial_balances.clone(),
84            current_coins: account.coins.clone(),
85            current_objects: vec![],
86            initial_data: account,
87        }
88    }
89
90    // TODO: Use this to get around the fact that we need to update object refs in the
91    // executor..figure out a better way to do this other than just creating a gas object for each
92    // transaction.
93    pub fn new_gas_object(&mut self, exec: &mut Executor) -> Object {
94        // We just create a new gas object for this transaction
95        let gas_object_id = ObjectID::random();
96        let gas_object = Object::with_id_owner_gas_for_testing(
97            gas_object_id,
98            self.initial_data.account.address,
99            self.initial_data.balance_creation_amt,
100        );
101        exec.add_object(gas_object.clone());
102        self.current_balances
103            .push(self.initial_data.balance_creation_amt);
104        self.current_coins.push(gas_object.clone());
105        gas_object
106    }
107}
108
109impl Arbitrary for Account {
110    type Parameters = ();
111    type Strategy = fn() -> Account;
112    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
113        Account::new_random as Self::Strategy
114    }
115}
116
117impl AccountData {
118    /// Returns a [`Strategy`] that creates `AccountData` instances.
119    pub fn strategy(balance_strategy: impl Strategy<Value = u64>) -> impl Strategy<Value = Self> {
120        (any::<Account>(), balance_strategy).prop_map(|(account, balance)| {
121            AccountData::new_with_account_and_balance(Arc::new(account), balance)
122        })
123    }
124}
transaction_fuzzer/account_universe/account.rs

transaction_fuzzer/account_universe/
account.rs
