icrc1_test_env_replica/
lib.rs

1use anyhow::Context;
2use async_trait::async_trait;
3use candid::utils::{decode_args, encode_args, ArgumentDecoder, ArgumentEncoder};
4use candid::Principal;
5use ic_agent::identity::BasicIdentity;
6use ic_agent::Agent;
7use icrc1_test_env::LedgerEnv;
8use ring::rand::SystemRandom;
9use std::sync::{Arc, Mutex};
10use std::time::SystemTime;
11
12pub fn fresh_identity(rand: &SystemRandom) -> BasicIdentity {
13    use ring::signature::Ed25519KeyPair as KeyPair;
14
15    let doc = KeyPair::generate_pkcs8(rand).expect("failed to generate an ed25519 key pair");
16
17    let key_pair = KeyPair::from_pkcs8(doc.as_ref())
18        .expect("failed to construct a key pair from a pkcs8 document");
19    BasicIdentity::from_key_pair(key_pair)
20}
21
22#[derive(Clone)]
23pub struct ReplicaLedger {
24    rand: Arc<Mutex<SystemRandom>>,
25    agent: Arc<Agent>,
26    canister_id: Principal,
27}
28
29#[async_trait(?Send)]
30impl LedgerEnv for ReplicaLedger {
31    fn fork(&self) -> Self {
32        let mut agent = Arc::clone(&self.agent);
33        Arc::make_mut(&mut agent).set_identity({
34            let r = self.rand.lock().expect("failed to grab a lock");
35            fresh_identity(&r)
36        });
37        Self {
38            rand: Arc::clone(&self.rand),
39            agent,
40            canister_id: self.canister_id,
41        }
42    }
43
44    fn principal(&self) -> Principal {
45        self.agent
46            .get_principal()
47            .expect("failed to get agent principal")
48    }
49
50    async fn time(&self) -> SystemTime {
51        // The replica relies on the system time by default.
52        // Unfortunately, this assumption might break during the time
53        // shifts, but it's probably good enough for tests.
54        SystemTime::now()
55    }
56
57    async fn query<Input, Output>(&self, method: &str, input: Input) -> anyhow::Result<Output>
58    where
59        Input: ArgumentEncoder + std::fmt::Debug,
60        Output: for<'a> ArgumentDecoder<'a>,
61    {
62        let debug_inputs = format!("{:?}", input);
63        let in_bytes = encode_args(input)
64            .with_context(|| format!("Failed to encode arguments {}", debug_inputs))?;
65        let bytes = self
66            .agent
67            .query(&self.canister_id, method)
68            .with_arg(in_bytes)
69            .call()
70            .await
71            .with_context(|| {
72                format!(
73                    "failed to call method {} on {} with args {}",
74                    method, self.canister_id, debug_inputs
75                )
76            })?;
77
78        decode_args(&bytes).with_context(|| {
79            format!(
80                "Failed to decode method {} response into type {}, bytes: {}",
81                method,
82                std::any::type_name::<Output>(),
83                hex::encode(bytes)
84            )
85        })
86    }
87
88    async fn update<Input, Output>(&self, method: &str, input: Input) -> anyhow::Result<Output>
89    where
90        Input: ArgumentEncoder + std::fmt::Debug,
91        Output: for<'a> ArgumentDecoder<'a>,
92    {
93        let debug_inputs = format!("{:?}", input);
94        let in_bytes = encode_args(input)
95            .with_context(|| format!("Failed to encode arguments {}", debug_inputs))?;
96        let bytes = self
97            .agent
98            .update(&self.canister_id, method)
99            .with_arg(in_bytes)
100            .call_and_wait()
101            .await
102            .with_context(|| {
103                format!(
104                    "failed to call method {} on {} with args {}",
105                    method, self.canister_id, debug_inputs
106                )
107            })?;
108
109        decode_args(&bytes).with_context(|| {
110            format!(
111                "Failed to decode method {} response into type {}, bytes: {}",
112                method,
113                std::any::type_name::<Output>(),
114                hex::encode(bytes)
115            )
116        })
117    }
118}
119
120impl ReplicaLedger {
121    pub fn new(agent: Agent, canister_id: Principal) -> Self {
122        Self {
123            rand: Arc::new(Mutex::new(SystemRandom::new())),
124            agent: Arc::new(agent),
125            canister_id,
126        }
127    }
128}
icrc1_test_env_replica/lib.rs

icrc1_test_env_replica/
lib.rs
