icrc1_test_env_pocket_ic/
lib.rs

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

icrc1_test_env_pocket_ic/
lib.rs
