Skip to main content

surfpool_sdk/cheatcodes/builders/
set_token_account.rs

1use solana_pubkey::Pubkey;
2
3use crate::cheatcodes::builders::CheatcodeBuilder;
4
5/// Builder for `surfnet_setTokenAccount`.
6///
7/// The required inputs are the token owner wallet and mint. Optional methods
8/// can then be used to set token-account fields before execution.
9///
10/// ```rust
11/// use surfpool_sdk::{Pubkey, Surfnet};
12/// use surfpool_sdk::cheatcodes::builders::SetTokenAccount;
13///
14/// # async fn example() {
15/// let surfnet = Surfnet::start().await.unwrap();
16/// let cheats = surfnet.cheatcodes();
17/// let owner = Pubkey::new_unique();
18/// let mint = Pubkey::new_unique();
19///
20/// cheats
21///     .execute(SetTokenAccount::new(owner, mint).amount(1_000))
22///     .unwrap();
23/// # }
24/// ```
25pub struct SetTokenAccount {
26    owner: Pubkey,
27    mint: Pubkey,
28    amount: Option<u64>,
29    delegate: Option<Option<Pubkey>>,
30    state: Option<String>,
31    delegated_amount: Option<u64>,
32    close_authority: Option<Option<Pubkey>>,
33    token_program: Option<Pubkey>,
34}
35
36impl SetTokenAccount {
37    /// Create a new token-account update builder for the given owner and mint.
38    pub fn new(owner: Pubkey, mint: Pubkey) -> Self {
39        Self {
40            owner,
41            mint,
42            amount: None,
43            delegate: None,
44            state: None,
45            delegated_amount: None,
46            close_authority: None,
47            token_program: None,
48        }
49    }
50
51    /// Set the token amount for the associated token account.
52    pub fn amount(mut self, amount: u64) -> Self {
53        self.amount = Some(amount);
54        self
55    }
56
57    /// Set a delegate authority on the token account.
58    pub fn delegate(mut self, delegate: Pubkey) -> Self {
59        self.delegate = Some(Some(delegate));
60        self
61    }
62
63    /// Clear any existing delegate authority.
64    pub fn clear_delegate(mut self) -> Self {
65        self.delegate = Some(None);
66        self
67    }
68
69    /// Set the token account state string expected by Surfnet RPC.
70    pub fn state(mut self, state: impl Into<String>) -> Self {
71        self.state = Some(state.into());
72        self
73    }
74
75    /// Set the delegated token amount.
76    pub fn delegated_amount(mut self, delegated_amount: u64) -> Self {
77        self.delegated_amount = Some(delegated_amount);
78        self
79    }
80
81    /// Set the close authority on the token account.
82    pub fn close_authority(mut self, close_authority: Pubkey) -> Self {
83        self.close_authority = Some(Some(close_authority));
84        self
85    }
86
87    /// Clear any existing close authority.
88    pub fn clear_close_authority(mut self) -> Self {
89        self.close_authority = Some(None);
90        self
91    }
92
93    /// Override the token program id.
94    ///
95    /// If omitted, the classic SPL Token program is used.
96    pub fn token_program(mut self, token_program: Pubkey) -> Self {
97        self.token_program = Some(token_program);
98        self
99    }
100}
101
102impl CheatcodeBuilder for SetTokenAccount {
103    const METHOD: &'static str = "surfnet_setTokenAccount";
104
105    /// Build the JSON-RPC parameter array for `surfnet_setTokenAccount`.
106    fn build(self) -> serde_json::Value {
107        let mut update = serde_json::Map::new();
108        if let Some(amount) = self.amount {
109            update.insert("amount".to_string(), amount.into());
110        }
111        if let Some(delegate) = self.delegate {
112            update.insert(
113                "delegate".to_string(),
114                delegate
115                    .map(|pubkey| pubkey.to_string().into())
116                    .unwrap_or_else(|| "null".into()),
117            );
118        }
119        if let Some(state) = self.state {
120            update.insert("state".to_string(), state.into());
121        }
122        if let Some(delegated_amount) = self.delegated_amount {
123            update.insert("delegatedAmount".to_string(), delegated_amount.into());
124        }
125        if let Some(close_authority) = self.close_authority {
126            update.insert(
127                "closeAuthority".to_string(),
128                close_authority
129                    .map(|pubkey| pubkey.to_string().into())
130                    .unwrap_or_else(|| "null".into()),
131            );
132        }
133
134        let mut params = vec![
135            self.owner.to_string().into(),
136            self.mint.to_string().into(),
137            update.into(),
138        ];
139
140        if let Some(token_program) = self.token_program {
141            params.push(token_program.to_string().into());
142        }
143
144        serde_json::Value::Array(params)
145    }
146}