sui_cluster_test/
helper.rs1use 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#[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#[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 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}