Skip to main content

sui_cluster_test/
helper.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::bail;
5use sui_rpc_api::Client as GrpcClient;
6use sui_sdk_types::BalanceChange;
7use sui_types::base_types::SuiAddress;
8use sui_types::gas_coin::GasCoin;
9use sui_types::object::Object;
10use sui_types::sui_sdk_types_conversions::type_tag_core_to_sdk;
11use sui_types::{base_types::ObjectID, object::Owner, parse_sui_type_tag};
12use tracing::{debug, trace};
13
14/// A util struct that helps verify a Sui object over gRPC (`LedgerService`).
15/// Use builder style to construct the conditions. When optional fields are not
16/// set, related checks are omitted. Consuming functions such as `check` perform
17/// the check and panic if verification results are unexpected. `check_into_object`
18/// and `check_into_gas_coin` return the native `Object` and `GasCoin`
19/// respectively.
20///
21/// Deleted/wrapped/unwrapped dispositions are verified through transaction
22/// effects at the call sites, not by observing a `get_object` failure, so this
23/// checker only handles the "object still exists" reads.
24#[derive(Debug)]
25pub struct ObjectChecker {
26    object_id: ObjectID,
27    owner: Option<Owner>,
28    is_sui_coin: Option<bool>,
29}
30
31impl ObjectChecker {
32    pub fn new(object_id: ObjectID) -> ObjectChecker {
33        Self {
34            object_id,
35            owner: None,
36            is_sui_coin: None,
37        }
38    }
39
40    pub fn owner(mut self, owner: Owner) -> Self {
41        self.owner = Some(owner);
42        self
43    }
44
45    pub fn is_sui_coin(mut self, is_sui_coin: bool) -> Self {
46        self.is_sui_coin = Some(is_sui_coin);
47        self
48    }
49
50    pub async fn check_into_gas_coin(self, client: &GrpcClient) -> GasCoin {
51        if self.is_sui_coin == Some(false) {
52            panic!("'check_into_gas_coin' shouldn't be called with 'is_sui_coin' set as false");
53        }
54        self.is_sui_coin(true)
55            .check(client)
56            .await
57            .unwrap()
58            .into_gas_coin()
59    }
60
61    pub async fn check_into_object(self, client: &GrpcClient) -> Object {
62        self.check(client).await.unwrap().into_object()
63    }
64
65    pub async fn check(self, client: &GrpcClient) -> Result<CheckerResultObject, anyhow::Error> {
66        debug!(?self);
67
68        let object_id = self.object_id;
69        let mut client = client.clone();
70        let object = match client.get_object(object_id).await {
71            Ok(object) => object,
72            Err(err) => bail!("Failed to get object info (id: {}), err: {err}", object_id),
73        };
74
75        trace!("getting object {object_id}, info :: {object:?}");
76
77        if let Some(owner) = self.owner {
78            let object_owner = object.owner().clone();
79            assert_eq!(
80                object_owner, owner,
81                "Object {} does not belong to {}, but {}",
82                object_id, owner, object_owner
83            );
84        }
85
86        if self.is_sui_coin == Some(true) {
87            let gas_coin = GasCoin::try_from(&object)
88                .map_err(|e| anyhow::anyhow!("Object {} is not a SUI gas coin: {e}", object_id))?;
89            return Ok(CheckerResultObject::new(Some(gas_coin), Some(object)));
90        }
91
92        Ok(CheckerResultObject::new(None, Some(object)))
93    }
94}
95
96pub struct CheckerResultObject {
97    gas_coin: Option<GasCoin>,
98    object: Option<Object>,
99}
100
101impl CheckerResultObject {
102    pub fn new(gas_coin: Option<GasCoin>, object: Option<Object>) -> Self {
103        Self { gas_coin, object }
104    }
105    pub fn into_gas_coin(self) -> GasCoin {
106        self.gas_coin.unwrap()
107    }
108    pub fn into_object(self) -> Object {
109        self.object.unwrap()
110    }
111}
112
113#[macro_export]
114macro_rules! assert_eq_if_present {
115    ($left:expr, $right:expr, $($arg:tt)+) => {
116        match (&$left, &$right) {
117            (Some(left_val), right_val) if !(&left_val == right_val) => {
118                panic!("{} does not match, left: {:?}, right: {:?}", $($arg)+, left_val, right_val);
119            }
120            _ => ()
121        }
122    };
123}
124
125/// Verifies a native SDK balance change (`sui_sdk_types::BalanceChange`) returned
126/// by the gRPC execution result. Coin types are compared exactly (as canonical
127/// `TypeTag`s), never by substring matching.
128#[derive(Default, Debug)]
129pub struct BalanceChangeChecker {
130    address: Option<SuiAddress>,
131    coin_type: Option<sui_sdk_types::TypeTag>,
132    amount: Option<i128>,
133}
134
135impl BalanceChangeChecker {
136    pub fn new() -> Self {
137        Default::default()
138    }
139
140    pub fn address(mut self, address: SuiAddress) -> Self {
141        self.address = Some(address);
142        self
143    }
144
145    pub fn coin_type(mut self, coin_type: &str) -> Self {
146        // Parse into a native `TypeTag` then into the SDK type so we do an exact,
147        // canonical comparison rather than a string/substring match.
148        let type_tag = parse_sui_type_tag(coin_type).unwrap();
149        let sdk_type_tag =
150            type_tag_core_to_sdk(type_tag).expect("coin type should convert into an SDK TypeTag");
151        self.coin_type = Some(sdk_type_tag);
152        self
153    }
154
155    pub fn amount(mut self, amount: i128) -> Self {
156        self.amount = Some(amount);
157        self
158    }
159
160    pub fn check(self, change: &BalanceChange) {
161        let BalanceChange {
162            address,
163            coin_type,
164            amount,
165        } = change;
166
167        if let Some(expected) = self.address {
168            let expected_sdk: sui_sdk_types::Address = expected.into();
169            assert_eq!(
170                &expected_sdk, address,
171                "balance change address does not match"
172            );
173        }
174        assert_eq_if_present!(self.coin_type, coin_type, "coin_type");
175        assert_eq_if_present!(self.amount, amount, "amount");
176    }
177}