icrc1_test_env_state_machine/
lib.rs

1use anyhow::Context;
2use async_trait::async_trait;
3use candid::utils::{decode_args, encode_args, ArgumentDecoder, ArgumentEncoder};
4use candid::Principal;
5use ic_test_state_machine_client::StateMachine;
6use icrc1_test_env::LedgerEnv;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::Arc;
9
10fn new_principal(n: u64) -> Principal {
11    let mut bytes = n.to_le_bytes().to_vec();
12    bytes.push(0xfe);
13    bytes.push(0x01);
14    Principal::try_from_slice(&bytes[..]).unwrap()
15}
16
17#[derive(Clone)]
18pub struct SMLedger {
19    counter: Arc<AtomicU64>,
20    sm: Arc<StateMachine>,
21    sender: Principal,
22    canister_id: Principal,
23}
24
25#[async_trait(?Send)]
26impl LedgerEnv for SMLedger {
27    fn fork(&self) -> Self {
28        Self {
29            counter: self.counter.clone(),
30            sm: self.sm.clone(),
31            sender: new_principal(self.counter.fetch_add(1, Ordering::Relaxed)),
32            canister_id: self.canister_id,
33        }
34    }
35
36    fn principal(&self) -> Principal {
37        self.sender
38    }
39
40    fn time(&self) -> std::time::SystemTime {
41        self.sm.time()
42    }
43
44    async fn query<Input, Output>(&self, method: &str, input: Input) -> anyhow::Result<Output>
45    where
46        Input: ArgumentEncoder + std::fmt::Debug,
47        Output: for<'a> ArgumentDecoder<'a>,
48    {
49        let debug_inputs = format!("{:?}", input);
50        let in_bytes = encode_args(input)
51            .with_context(|| format!("Failed to encode arguments {}", debug_inputs))?;
52        match self
53            .sm
54            .query_call(
55                Principal::from_slice(self.canister_id.as_slice()),
56                Principal::from_slice(self.sender.as_slice()),
57                method,
58                in_bytes,
59            )
60            .map_err(|err| anyhow::Error::msg(err.to_string()))?
61        {
62            ic_test_state_machine_client::WasmResult::Reply(bytes) => decode_args(&bytes)
63                .with_context(|| {
64                    format!(
65                        "Failed to decode method {} response into type {}, bytes: {}",
66                        method,
67                        std::any::type_name::<Output>(),
68                        hex::encode(bytes)
69                    )
70                }),
71            ic_test_state_machine_client::WasmResult::Reject(msg) => {
72                return Err(anyhow::Error::msg(format!(
73                    "Query call to ledger {:?} was rejected: {}",
74                    self.canister_id, msg
75                )))
76            }
77        }
78    }
79
80    async fn update<Input, Output>(&self, method: &str, input: Input) -> anyhow::Result<Output>
81    where
82        Input: ArgumentEncoder + std::fmt::Debug,
83        Output: for<'a> ArgumentDecoder<'a>,
84    {
85        let debug_inputs = format!("{:?}", input);
86        let in_bytes = encode_args(input)
87            .with_context(|| format!("Failed to encode arguments {}", debug_inputs))?;
88        match self
89            .sm
90            .update_call(self.canister_id, self.sender, method, in_bytes)
91            .map_err(|err| anyhow::Error::msg(err.to_string()))
92            .with_context(|| {
93                format!(
94                    "failed to execute update call {} on canister {}",
95                    method, self.canister_id
96                )
97            })? {
98            ic_test_state_machine_client::WasmResult::Reply(bytes) => decode_args(&bytes)
99                .with_context(|| {
100                    format!(
101                        "Failed to decode method {} response into type {}, bytes: {}",
102                        method,
103                        std::any::type_name::<Output>(),
104                        hex::encode(bytes)
105                    )
106                }),
107            ic_test_state_machine_client::WasmResult::Reject(msg) => {
108                return Err(anyhow::Error::msg(format!(
109                    "Query call to ledger {:?} was rejected: {}",
110                    self.canister_id, msg
111                )))
112            }
113        }
114    }
115}
116
117impl SMLedger {
118    pub fn new(sm: Arc<StateMachine>, canister_id: Principal, sender: Principal) -> Self {
119        Self {
120            counter: Arc::new(AtomicU64::new(0)),
121            sm,
122            canister_id,
123            sender,
124        }
125    }
126}
icrc1_test_env_state_machine/lib.rs

icrc1_test_env_state_machine/
lib.rs
