Skip to main content

ic_testkit/pocket_ic/
nonblocking.rs

1// Copyright DFINITY Foundation. Licensed under Apache-2.0.
2use crate::pocket_ic::common::rest::{
3    AutoProgressConfig, BlobCompression, BlobId, CanisterHttpRequest, CreateHttpGatewayResponse,
4    CreateInstanceResponse, HttpGatewayBackend, HttpGatewayConfig, HttpGatewayInfo, HttpsConfig,
5    InstanceConfig, InstanceId, MockCanisterHttpResponse, RawAddCycles, RawCanisterCall,
6    RawCanisterHttpRequest, RawCanisterId, RawCanisterResult, RawCanisterSnapshotDownload,
7    RawCanisterSnapshotId, RawCanisterSnapshotUpload, RawCycles, RawEffectivePrincipal,
8    RawIngressStatusArgs, RawMessageId, RawMockCanisterHttpResponse, RawPrincipalId,
9    RawSetStableMemory, RawStableMemory, RawSubnetId, RawTickConfigs, RawTime, SubnetId, Topology,
10};
11#[cfg(windows)]
12use crate::pocket_ic::wsl_path;
13use crate::pocket_ic::{ClientError, Operation, RequestContext, http::Exchange};
14use crate::pocket_ic::{
15    CreateCanisterParams, CreateCanisterPlacement, IngressStatusResult, PocketIcBuilder,
16    RejectResponse, TickConfigs, Time,
17};
18use backon::{BackoffBuilder, ExponentialBuilder};
19use candid::{
20    Principal, decode_args, encode_args,
21    utils::{ArgumentDecoder, ArgumentEncoder},
22};
23use ic_management_canister_types::{
24    CanisterId, CanisterIdRecord, CanisterInstallMode, CanisterLogRecord, CanisterSettings,
25    CanisterStatusResult, ChunkHash, DeleteCanisterSnapshotArgs, FetchCanisterLogsResult,
26    InstallChunkedCodeArgs, InstallCodeArgs, LoadCanisterSnapshotArgs,
27    ProvisionalCreateCanisterWithCyclesArgs, Snapshot, StoredChunksResult,
28    TakeCanisterSnapshotArgs, UpdateSettingsArgs, UploadChunkArgs, UploadChunkResult,
29};
30use ic_platform::certification::cbor::decode_certificate;
31use ic_platform::certification::{Certificate, Label, LookupResult};
32use ic_platform::ic_protocol::Envelope;
33use ic_platform::ic_protocol::EnvelopeContent::ReadState;
34use ic_platform::ic_protocol::{
35    ReadStateResponse, SubnetMetrics, decode_read_state_response, decode_subnet_metrics,
36    encode_envelope,
37};
38use reqwest::{StatusCode, Url};
39use serde::{Serialize, de::DeserializeOwned};
40use sha2::{Digest, Sha256};
41use std::future::Future;
42use std::net::IpAddr;
43use std::path::PathBuf;
44use std::time::Duration;
45use tracing::{debug, instrument};
46use tracing_appender::non_blocking::WorkerGuard;
47use tracing_subscriber::EnvFilter;
48
49// wait time between polling requests
50
51mod snapshot;
52
53const LOG_DIR_PATH_ENV_NAME: &str = "POCKET_IC_LOG_DIR";
54const LOG_DIR_LEVELS_ENV_NAME: &str = "POCKET_IC_LOG_DIR_LEVELS";
55
56const LOCALHOST: &str = "localhost";
57
58// The minimum joint size of a canister's WASM
59// and its initial argument blob for which
60// we install the canister WASM as a sequence of chunks.
61// The size is chosen to be smaller than the maximum
62// ingress message size of 2 MiB (2,097,152 B)
63// on application subnet (which have the tighest
64// ingress message size limit)
65// to account for a few additional constant-size fields
66// in the management canister payload to install code
67// and the overhead of candid encoding.
68const INSTALL_CHUNKED_CODE_THRESHOLD: usize = 2_000_000; // 2 MB
69
70// The maximum size of one WASM chunk when installing a canister WASM
71// as a sequence of chunks. This constant is specified
72// in the IC protocol.
73const INSTALL_CODE_CHUNK_SIZE: usize = 1 << 20; // 1 MiB
74
75enum HttpMethod {
76    Get,
77    Post,
78}
79
80/// Main entry point for interacting with PocketIC.
81pub struct PocketIc {
82    /// The unique ID of this PocketIC instance.
83    pub instance_id: InstanceId,
84    // how long a get/post request may retry or poll
85    max_request_time_ms: Option<u64>,
86    http_gateway: Option<HttpGatewayInfo>,
87    server_url: Url,
88    reqwest_client: reqwest::Client,
89    // the instance should only be deleted when dropping this handle if this handle owns the instance
90    owns_instance: bool,
91    _log_guard: Option<WorkerGuard>,
92}
93
94impl PocketIc {
95    /// Creates a new PocketIC instance with a single application subnet on the server.
96    /// The caller supplies an already admitted server endpoint.
97    pub async fn new(server_url: Url) -> Result<Self, ClientError> {
98        PocketIcBuilder::new()
99            .with_server_url(server_url)
100            .with_application_subnet()
101            .build()
102            .await
103    }
104
105    /// Creates a PocketIC handle to an existing instance on a running server.
106    /// Note that this handle does not extend the lifetime of the existing instance,
107    /// i.e., the existing instance is deleted and this handle stops working
108    /// when the PocketIC handle that created the existing instance is dropped.
109    pub fn new_from_existing_instance(
110        server_url: Url,
111        instance_id: InstanceId,
112        max_request_time_ms: Option<u64>,
113    ) -> Self {
114        let test_driver_pid = std::process::id();
115        let log_guard = setup_tracing(test_driver_pid);
116
117        let reqwest_client = reqwest::Client::new();
118        debug!("instance_id={} Reusing existing instance.", instance_id);
119
120        Self {
121            instance_id,
122            max_request_time_ms,
123            http_gateway: None,
124            server_url,
125            reqwest_client,
126            owns_instance: false,
127            _log_guard: log_guard,
128        }
129    }
130
131    pub(crate) async fn from_config(
132        server_url: Url,
133        max_request_time_ms: Option<u64>,
134        instance_config: InstanceConfig,
135    ) -> Result<Self, ClientError> {
136        #[cfg(windows)]
137        let instance_config = InstanceConfig {
138            state_dir: instance_config
139                .state_dir
140                .map(|path| wsl_path(&path, "state directory").into()),
141            ..instance_config
142        };
143
144        let test_driver_pid = std::process::id();
145        let log_guard = setup_tracing(test_driver_pid);
146
147        let reqwest_client = reqwest::Client::new();
148        let url = server_url.join("instances").unwrap();
149        let mut exchange = Exchange::new(&url, max_request_time_ms);
150        let response: CreateInstanceResponse = exchange
151            .json(
152                &reqwest_client,
153                reqwest_client.post(url.clone()).json(&instance_config),
154            )
155            .await?;
156        let (instance_id, http_gateway_info) = match response {
157            CreateInstanceResponse::Created {
158                instance_id,
159                http_gateway_info,
160                ..
161            } => (instance_id, http_gateway_info),
162            CreateInstanceResponse::Error { message } => {
163                return Err(ClientError::Server {
164                    context: Box::new(exchange.context.clone()),
165                    status: exchange.context.status.unwrap_or(StatusCode::OK),
166                    message,
167                });
168            }
169        };
170        debug!("instance_id={} New instance created.", instance_id);
171
172        Ok(Self {
173            instance_id,
174            max_request_time_ms,
175            http_gateway: http_gateway_info,
176            server_url,
177            reqwest_client,
178            owns_instance: true,
179            _log_guard: log_guard,
180        })
181    }
182
183    /// Returns the URL of the PocketIC server on which this PocketIC instance is running.
184    pub fn get_server_url(&self) -> Url {
185        self.server_url.clone()
186    }
187
188    /// Returns the topology of the different subnets of this PocketIC instance.
189    pub async fn topology(&self) -> Result<Topology, ClientError> {
190        let endpoint = "read/topology";
191        self.get(endpoint).await
192    }
193
194    /// Upload and store a binary blob to the PocketIC server.
195    #[instrument(ret, skip(self, blob), fields(instance_id=self.instance_id, blob_len = %blob.len(), compression = ?compression))]
196    pub async fn upload_blob(
197        &self,
198        blob: Vec<u8>,
199        compression: BlobCompression,
200    ) -> Result<BlobId, ClientError> {
201        let url = self.server_url.join("blobstore/").unwrap();
202        let mut request = self.reqwest_client.post(url.clone()).body(blob);
203        if let BlobCompression::Gzip = compression {
204            request = request.header(reqwest::header::CONTENT_ENCODING, "gzip");
205        }
206        let mut exchange = Exchange::new(&url, self.max_request_time_ms);
207        let (status, bytes) = exchange.bytes(&self.reqwest_client, request).await?;
208        if !status.is_success() {
209            return Err(exchange.server_error(status, &bytes));
210        }
211        let hash = hex::decode(bytes).map_err(|error| ClientError::Protocol(error.to_string()))?;
212        Ok(BlobId(hash))
213    }
214
215    /// Set stable memory of a canister. Optional GZIP compression can be used for reduced
216    /// data traffic.
217    #[instrument(skip(self, data), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), data_len = %data.len(), compression = ?compression))]
218    pub async fn set_stable_memory(
219        &self,
220        canister_id: CanisterId,
221        data: Vec<u8>,
222        compression: BlobCompression,
223    ) -> Result<(), ClientError> {
224        let blob_id = self.upload_blob(data, compression).await?;
225        let endpoint = "update/set_stable_memory";
226        self.post::<(), _>(
227            endpoint,
228            RawSetStableMemory {
229                canister_id: canister_id.as_slice().to_vec(),
230                blob_id,
231            },
232        )
233        .await?;
234
235        Ok(())
236    }
237
238    /// Get stable memory of a canister.
239    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
240    pub async fn get_stable_memory(&self, canister_id: CanisterId) -> Result<Vec<u8>, ClientError> {
241        let endpoint = "read/get_stable_memory";
242        let RawStableMemory { blob } = self
243            .post(
244                endpoint,
245                RawCanisterId {
246                    canister_id: canister_id.as_slice().to_vec(),
247                },
248            )
249            .await?;
250        Ok(blob)
251    }
252
253    /// List all instances and their status.
254    #[instrument(ret)]
255    pub async fn list_instances(server_url: Url) -> Result<Vec<String>, ClientError> {
256        let url = server_url.join("instances").unwrap();
257        let client = reqwest::Client::new();
258        Exchange::new(&url, None)
259            .json(&client, client.get(url.clone()))
260            .await
261    }
262
263    /// Make the IC produce and progress by one block.
264    /// Note that multiple ticks might be necessary to observe
265    /// an expected effect, e.g., if the effect depends on
266    /// inter-canister calls or heartbeats.
267    #[instrument(skip(self), fields(instance_id=self.instance_id))]
268    pub async fn tick(&self) -> Result<(), ClientError> {
269        self.tick_with_configs(TickConfigs::default()).await?;
270
271        Ok(())
272    }
273
274    /// Make the IC produce and progress by one block with custom
275    /// configs for the round.
276    #[instrument(skip(self), fields(instance_id=self.instance_id))]
277    pub async fn tick_with_configs(&self, configs: TickConfigs) -> Result<(), ClientError> {
278        let endpoint = "update/tick";
279        self.post::<(), RawTickConfigs>(endpoint, configs.into())
280            .await?;
281
282        Ok(())
283    }
284
285    /// Configures the IC to make progress automatically,
286    /// i.e., periodically update the time of the IC
287    /// to the real time and execute rounds on the subnets.
288    /// Only returns after the certified time of the IC
289    /// has been updated for the first time.
290    /// Returns the URL at which `/api` requests
291    /// for this instance can be made.
292    #[instrument(skip(self), fields(instance_id=self.instance_id))]
293    pub async fn auto_progress(&self) -> Result<Url, ClientError> {
294        let endpoint = "auto_progress";
295        let auto_progress_config = AutoProgressConfig {
296            artificial_delay_ms: None,
297        };
298        self.post::<(), _>(endpoint, auto_progress_config).await?;
299        Ok(self.instance_url())
300    }
301
302    /// Returns whether automatic progress is enabled on the PocketIC instance.
303    #[instrument(skip(self), fields(instance_id=self.instance_id))]
304    pub async fn auto_progress_enabled(&self) -> Result<bool, ClientError> {
305        self.get("auto_progress").await
306    }
307
308    pub(crate) fn instance_url(&self) -> Url {
309        self.server_url
310            .join("/instances/")
311            .unwrap()
312            .join(&format!("{}/", self.instance_id))
313            .unwrap()
314    }
315
316    /// Stops automatic progress (see `auto_progress`) on the IC.
317    #[instrument(skip(self), fields(instance_id=self.instance_id))]
318    pub async fn stop_progress(&self) -> Result<(), ClientError> {
319        let endpoint = "stop_progress";
320        self.post::<(), _>(endpoint, "").await?;
321
322        Ok(())
323    }
324
325    /// Returns the URL at which `/api` requests
326    /// for this instance can be made if the HTTP
327    /// gateway has been started.
328    pub fn url(&self) -> Option<Url> {
329        self.http_gateway
330            .as_ref()
331            .map(|res| Url::parse(&format!("http://{}:{}/", LOCALHOST, res.port)).unwrap())
332    }
333
334    /// Creates an HTTP gateway for this PocketIC instance binding to `127.0.0.1`
335    /// and an optionally specified port (defaults to choosing an arbitrary unassigned port);
336    /// listening on `localhost`;
337    /// and configures the PocketIC instance to make progress automatically, i.e.,
338    /// periodically update the time of the PocketIC instance to the real time
339    /// and process messages on the PocketIC instance.
340    /// Only returns after the certified time of the PocketIC instance
341    /// has been updated for the first time.
342    /// Returns the URL at which `/api` requests
343    /// for this instance can be made.
344    #[instrument(skip(self), fields(instance_id=self.instance_id))]
345    pub async fn make_live(&mut self, listen_at: Option<u16>) -> Result<Url, ClientError> {
346        self.make_live_with_params(None, listen_at, None, None)
347            .await
348    }
349
350    /// Creates an HTTP gateway for this PocketIC instance binding
351    /// to an optionally specified IP address (defaults to `127.0.0.1`)
352    /// and port (defaults to choosing an arbitrary unassigned port);
353    /// listening on optionally specified domains (default to `localhost`);
354    /// and using an optionally specified TLS certificate (if provided, an HTTPS gateway is created)
355    /// and configures the PocketIC instance to make progress automatically, i.e.,
356    /// periodically update the time of the PocketIC instance to the real time
357    /// and process messages on the PocketIC instance.
358    /// Only returns after the certified time of the PocketIC instance
359    /// has been updated for the first time.
360    /// Returns the URL at which `/api` requests
361    /// for this instance can be made.
362    #[instrument(skip(self), fields(instance_id=self.instance_id))]
363    pub async fn make_live_with_params(
364        &mut self,
365        ip_addr: Option<IpAddr>,
366        listen_at: Option<u16>,
367        domains: Option<Vec<String>>,
368        https_config: Option<HttpsConfig>,
369    ) -> Result<Url, ClientError> {
370        if let Some(url) = self.url() {
371            return Ok(url);
372        }
373        if !self.auto_progress_enabled().await? {
374            self.auto_progress().await?;
375        }
376        self.start_http_gateway(
377            ip_addr.map(|ip_addr| ip_addr.to_string()),
378            listen_at,
379            domains,
380            https_config,
381        )
382        .await
383    }
384
385    async fn start_http_gateway(
386        &mut self,
387        ip_addr: Option<String>,
388        port: Option<u16>,
389        domains: Option<Vec<String>>,
390        https_config: Option<HttpsConfig>,
391    ) -> Result<Url, ClientError> {
392        let endpoint = self.server_url.join("http_gateway").unwrap();
393        let http_gateway_config = HttpGatewayConfig {
394            ip_addr,
395            port,
396            forward_to: HttpGatewayBackend::PocketIcInstance(self.instance_id),
397            domains: domains.clone(),
398            https_config: https_config.clone(),
399            domain_custom_provider_local_file: None,
400        };
401        let mut exchange = Exchange::new(&endpoint, self.max_request_time_ms);
402        let res: CreateHttpGatewayResponse = exchange
403            .json(
404                &self.reqwest_client,
405                self.reqwest_client
406                    .post(endpoint.clone())
407                    .json(&http_gateway_config),
408            )
409            .await?;
410        Ok(match res {
411            CreateHttpGatewayResponse::Created(info) => {
412                let port = info.port;
413                self.http_gateway = Some(info);
414                let proto = if https_config.is_some() {
415                    "https"
416                } else {
417                    "http"
418                };
419                Url::parse(&format!(
420                    "{}://{}:{}/",
421                    proto,
422                    domains
423                        .unwrap_or_default()
424                        .into_iter()
425                        .next()
426                        .unwrap_or(LOCALHOST.to_string()),
427                    port
428                ))
429                .unwrap()
430            }
431            CreateHttpGatewayResponse::Error { message } => {
432                return Err(ClientError::Server {
433                    context: Box::new(exchange.context.clone()),
434                    status: exchange.context.status.unwrap_or(StatusCode::OK),
435                    message,
436                });
437            }
438        })
439    }
440
441    async fn stop_http_gateway(&mut self) -> Result<(), ClientError> {
442        if let Some(info) = &self.http_gateway {
443            let url = self
444                .server_url
445                .join(&format!("http_gateway/{}/stop", info.instance_id))
446                .unwrap();
447            Exchange::new(&url, self.max_request_time_ms)
448                .json::<()>(&self.reqwest_client, self.reqwest_client.post(url.clone()))
449                .await?;
450            self.http_gateway = None;
451        }
452        Ok(())
453    }
454
455    /// Stops auto progress (automatic time updates and round executions)
456    /// and the HTTP gateway for this IC instance.
457    #[instrument(skip(self), fields(instance_id=self.instance_id))]
458    pub async fn stop_live(&mut self) -> Result<(), ClientError> {
459        self.stop_http_gateway().await?;
460        self.stop_progress().await?;
461
462        Ok(())
463    }
464
465    /// Get the root key of this IC instance. Returns `None` if the IC has no NNS subnet.
466    #[instrument(skip(self), fields(instance_id=self.instance_id))]
467    pub async fn root_key(&self) -> Result<Option<Vec<u8>>, ClientError> {
468        let Some(subnet_id) = self.topology().await?.get_nns() else {
469            return Ok(None);
470        };
471        let subnet_id: RawSubnetId = subnet_id.into();
472        let endpoint = "read/pub_key";
473        let res = self.post::<Vec<u8>, _>(endpoint, subnet_id).await?;
474        Ok(Some(res))
475    }
476
477    /// Get the current time of the IC.
478    #[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
479    pub async fn get_time(&self) -> Result<Time, ClientError> {
480        let endpoint = "read/get_time";
481        let result: RawTime = self.get(endpoint).await?;
482        Ok(Time::from_nanos_since_unix_epoch(result.nanos_since_epoch))
483    }
484
485    /// Set the current time of the IC, on all subnets.
486    #[instrument(skip(self), fields(instance_id=self.instance_id, time = ?time))]
487    pub async fn set_time(&self, time: Time) -> Result<(), ClientError> {
488        let endpoint = "update/set_time";
489        self.post::<(), _>(
490            endpoint,
491            RawTime {
492                nanos_since_epoch: time.as_nanos_since_unix_epoch(),
493            },
494        )
495        .await?;
496
497        Ok(())
498    }
499
500    /// Set the current certified time of the IC, on all subnets.
501    #[instrument(skip(self), fields(instance_id=self.instance_id, time = ?time))]
502    pub async fn set_certified_time(&self, time: Time) -> Result<(), ClientError> {
503        let endpoint = "update/set_certified_time";
504        self.post::<(), _>(
505            endpoint,
506            RawTime {
507                nanos_since_epoch: time.as_nanos_since_unix_epoch(),
508            },
509        )
510        .await?;
511
512        Ok(())
513    }
514
515    /// Advance the time on the IC on all subnets by some nanoseconds.
516    #[instrument(skip(self), fields(instance_id=self.instance_id, duration = ?duration))]
517    pub async fn advance_time(&self, duration: Duration) -> Result<(), ClientError> {
518        let now = self.get_time().await?;
519        self.set_time(now + duration).await?;
520
521        Ok(())
522    }
523
524    /// Get the controllers of a canister.
525    /// Returns a server error if the canister does not exist.
526    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
527    pub async fn get_controllers(
528        &self,
529        canister_id: CanisterId,
530    ) -> Result<Vec<Principal>, ClientError> {
531        let endpoint = "read/get_controllers";
532        let result: Vec<RawPrincipalId> = self
533            .post(
534                endpoint,
535                RawCanisterId {
536                    canister_id: canister_id.as_slice().to_vec(),
537                },
538            )
539            .await?;
540        Ok(result.into_iter().map(|p| p.into()).collect())
541    }
542
543    /// Get the current cycles balance of a canister.
544    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
545    pub async fn cycle_balance(&self, canister_id: CanisterId) -> Result<u128, ClientError> {
546        let endpoint = "read/get_cycles";
547        let result: RawCycles = self
548            .post(
549                endpoint,
550                RawCanisterId {
551                    canister_id: canister_id.as_slice().to_vec(),
552                },
553            )
554            .await?;
555        Ok(result.cycles)
556    }
557
558    /// Add cycles to a canister. Returns the new balance.
559    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), amount = %amount))]
560    pub async fn add_cycles(
561        &self,
562        canister_id: CanisterId,
563        amount: u128,
564    ) -> Result<u128, ClientError> {
565        let endpoint = "update/add_cycles";
566        let result: RawCycles = self
567            .post(
568                endpoint,
569                RawAddCycles {
570                    canister_id: canister_id.as_slice().to_vec(),
571                    amount,
572                },
573            )
574            .await?;
575        Ok(result.cycles)
576    }
577
578    /// Submit an update call (without executing it immediately).
579    pub async fn submit_call(
580        &self,
581        canister_id: CanisterId,
582        sender: Principal,
583        method: &str,
584        payload: Vec<u8>,
585    ) -> Result<RawMessageId, ClientError> {
586        self.submit_call_with_effective_principal(
587            canister_id,
588            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
589            sender,
590            method,
591            payload,
592        )
593        .await
594    }
595
596    /// Submit an update with explicit routing, retaining its accepted ingress receipt.
597    pub async fn submit_call_with_effective_principal(
598        &self,
599        canister_id: CanisterId,
600        effective_principal: RawEffectivePrincipal,
601        sender: Principal,
602        method: &str,
603        payload: Vec<u8>,
604    ) -> Result<RawMessageId, ClientError> {
605        let endpoint = "update/submit_ingress_message";
606        let raw_canister_call = RawCanisterCall {
607            sender: sender.as_slice().to_vec(),
608            canister_id: canister_id.as_slice().to_vec(),
609            method: method.to_string(),
610            payload,
611            effective_principal,
612            sender_info: None,
613        };
614        self.post::<Result<RawMessageId, RejectResponse>, _>(endpoint, raw_canister_call)
615            .await?
616            .map_err(ClientError::from)
617    }
618
619    /// Await an update call submitted previously by `submit_call` or `submit_call_with_effective_principal`.
620    pub async fn await_call(&self, message_id: RawMessageId) -> Result<Vec<u8>, ClientError> {
621        let endpoint = "update/await_ingress_message";
622        let result: RawCanisterResult = self
623            .post(endpoint, &message_id)
624            .await
625            .map_err(|error| error.with_ingress(message_id))?;
626        Result::<Vec<u8>, RejectResponse>::from(result).map_err(ClientError::from)
627    }
628
629    /// Fetch the status of an update call submitted previously by `submit_call` or `submit_call_with_effective_principal`.
630    /// Note that the status of the update call can only change if the PocketIC instance is in live mode
631    /// or a round has been executed due to a separate PocketIC library call, e.g., `PocketIc::tick()`.
632    pub async fn ingress_status(
633        &self,
634        raw_message_id: RawMessageId,
635    ) -> Result<Option<Result<Vec<u8>, RejectResponse>>, ClientError> {
636        let status = self.ingress_status_as_caller(raw_message_id, None).await?;
637        Ok(match status {
638            IngressStatusResult::NotAvailable => None,
639            IngressStatusResult::Success(status) => Some(status),
640            IngressStatusResult::Forbidden(err) => return Err(ClientError::Protocol(err)),
641        })
642    }
643
644    /// Fetch the status of an update call submitted previously by `submit_call` or `submit_call_with_effective_principal`.
645    /// Note that the status of the update call can only change if the PocketIC instance is in live mode
646    /// or a round has been executed due to a separate PocketIC library call, e.g., `PocketIc::tick()`.
647    /// If the status of the update call is known, but the update call was submitted by a different caller, then an error is returned.
648    pub async fn ingress_status_as(
649        &self,
650        raw_message_id: RawMessageId,
651        caller: Principal,
652    ) -> Result<IngressStatusResult, ClientError> {
653        self.ingress_status_as_caller(raw_message_id, Some(caller))
654            .await
655    }
656
657    async fn ingress_status_as_caller(
658        &self,
659        raw_message_id: RawMessageId,
660        caller: Option<Principal>,
661    ) -> Result<IngressStatusResult, ClientError> {
662        let args = RawIngressStatusArgs {
663            raw_message_id: raw_message_id.clone(),
664            raw_caller: caller.map(Into::into),
665        };
666        match self
667            .post::<Option<RawCanisterResult>, _>("read/ingress_status", args)
668            .await
669        {
670            Ok(None) => Ok(IngressStatusResult::NotAvailable),
671            Ok(Some(result)) => Ok(IngressStatusResult::Success(result.into())),
672            Err(ClientError::Server {
673                status: StatusCode::FORBIDDEN,
674                message,
675                ..
676            }) => Ok(IngressStatusResult::Forbidden(message)),
677            Err(error) => Err(error.with_ingress(raw_message_id)),
678        }
679    }
680
681    /// Await an update call submitted previously by `submit_call` or `submit_call_with_effective_principal`.
682    /// Note that the status of the update call can only change if the PocketIC instance is in live mode
683    /// or a round has been executed due to a separate PocketIC library call.
684    pub async fn await_call_no_ticks(
685        &self,
686        message_id: RawMessageId,
687    ) -> Result<Vec<u8>, ClientError> {
688        let mut retry_policy = ExponentialBuilder::new()
689            .with_min_delay(Duration::from_millis(10))
690            .with_max_delay(Duration::from_secs(1))
691            .with_factor(2.0)
692            .with_jitter()
693            .with_total_delay(Some(Duration::from_secs(15 * 60)))
694            .without_max_times()
695            .build();
696        Ok(loop {
697            if let Some(ingress_status) = self.ingress_status(message_id.clone()).await? {
698                break ingress_status.map_err(ClientError::from)?;
699            }
700            let delay = retry_policy.next().ok_or_else(|| ClientError::Timeout {
701                context: Box::new(RequestContext {
702                    url: self.instance_url().to_string(),
703                    operation: None,
704                    ingress: Some(message_id.clone()),
705                    may_have_executed: true,
706                    status: None,
707                }),
708            })?;
709            tokio::time::sleep(delay).await;
710        })
711    }
712
713    /// Execute an update call on a canister.
714    #[instrument(skip(self, payload), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.to_string(), method = %method, payload_len = %payload.len()))]
715    pub async fn update_call(
716        &self,
717        canister_id: CanisterId,
718        sender: Principal,
719        method: &str,
720        payload: Vec<u8>,
721    ) -> Result<Vec<u8>, ClientError> {
722        self.update_call_with_effective_principal(
723            canister_id,
724            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
725            sender,
726            method,
727            payload,
728        )
729        .await
730    }
731
732    /// Submit once with explicit routing, then await the accepted ingress.
733    pub async fn update_call_with_effective_principal(
734        &self,
735        canister_id: CanisterId,
736        effective_principal: RawEffectivePrincipal,
737        sender: Principal,
738        method: &str,
739        payload: Vec<u8>,
740    ) -> Result<Vec<u8>, ClientError> {
741        let message_id = self
742            .submit_call_with_effective_principal(
743                canister_id,
744                effective_principal,
745                sender,
746                method,
747                payload,
748            )
749            .await?;
750        self.await_call(message_id).await
751    }
752
753    /// Execute a query call on a canister.
754    #[instrument(skip(self, payload), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.to_string(), method = %method, payload_len = %payload.len()))]
755    pub async fn query_call(
756        &self,
757        canister_id: CanisterId,
758        sender: Principal,
759        method: &str,
760        payload: Vec<u8>,
761    ) -> Result<Vec<u8>, ClientError> {
762        self.query_call_with_effective_principal(
763            canister_id,
764            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
765            sender,
766            method,
767            payload,
768        )
769        .await
770    }
771
772    /// Execute a query call on a canister explicitly specifying an effective principal to route the request:
773    /// this API is useful for making generic query calls (including management canister query calls) without using dedicated functions from this library
774    /// (e.g., making generic query calls in dfx to a PocketIC instance).
775    #[instrument(skip(self, payload), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), effective_principal = %effective_principal.to_string(), sender = %sender.to_string(), method = %method, payload_len = %payload.len()))]
776    pub async fn query_call_with_effective_principal(
777        &self,
778        canister_id: CanisterId,
779        effective_principal: RawEffectivePrincipal,
780        sender: Principal,
781        method: &str,
782        payload: Vec<u8>,
783    ) -> Result<Vec<u8>, ClientError> {
784        let endpoint = "read/query";
785        let raw_canister_call = RawCanisterCall {
786            sender: sender.as_slice().to_vec(),
787            canister_id: canister_id.as_slice().to_vec(),
788            method: method.to_string(),
789            payload,
790            effective_principal,
791            sender_info: None,
792        };
793        let result: RawCanisterResult = self.post(endpoint, raw_canister_call).await?;
794        Result::<Vec<u8>, RejectResponse>::from(result).map_err(ClientError::from)
795    }
796
797    /// Fetch canister logs via a query call to the management canister.
798    pub async fn fetch_canister_logs(
799        &self,
800        canister_id: CanisterId,
801        sender: Principal,
802    ) -> Result<Vec<CanisterLogRecord>, ClientError> {
803        Ok(with_candid::<_, (FetchCanisterLogsResult,), _>(
804            (CanisterIdRecord { canister_id },),
805            |payload| async {
806                self.query_call_with_effective_principal(
807                    Principal::management_canister(),
808                    RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
809                    sender,
810                    "fetch_canister_logs",
811                    payload,
812                )
813                .await
814            },
815        )
816        .await?
817        .0
818        .canister_log_records)
819    }
820
821    /// Request a canister's status.
822    #[instrument(skip(self), fields(instance_id=self.instance_id, sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
823    pub async fn canister_status(
824        &self,
825        canister_id: CanisterId,
826        sender: Option<Principal>,
827    ) -> Result<CanisterStatusResult, ClientError> {
828        Ok(
829            call_candid_as::<(CanisterIdRecord,), (CanisterStatusResult,)>(
830                self,
831                Principal::management_canister(),
832                RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
833                sender.unwrap_or(Principal::anonymous()),
834                "canister_status",
835                (CanisterIdRecord { canister_id },),
836            )
837            .await?
838            .0,
839        )
840    }
841
842    /// Create a canister with default settings as the anonymous principal.
843    /// The canister is created with 100T cycles.
844    #[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
845    pub async fn create_canister(&self) -> Result<CanisterId, ClientError> {
846        self.create_canister_with_params(None, CreateCanisterParams::default())
847            .await
848    }
849
850    /// Create a canister with optional custom settings and a sender.
851    /// The canister is created with 100T cycles.
852    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, settings = ?settings, sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
853    pub async fn create_canister_with_settings(
854        &self,
855        sender: Option<Principal>,
856        settings: Option<CanisterSettings>,
857    ) -> Result<CanisterId, ClientError> {
858        self.create_canister_with_params(
859            sender,
860            CreateCanisterParams {
861                settings,
862                ..Default::default()
863            },
864        )
865        .await
866    }
867
868    /// Creates a canister with a specific canister ID and optional custom settings.
869    /// The canister is created with 100T cycles.
870    /// Returns an error if the canister ID is already in use.
871    /// Creates a new subnet if the canister ID is not contained in any of the subnets.
872    ///
873    /// The canister ID must be an IC mainnet canister ID that does not belong to the NNS or II subnet,
874    /// otherwise the function might panic (for NNS and II canister IDs,
875    /// the PocketIC instance should already be created with those subnets).
876    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, sender = %sender.unwrap_or(Principal::anonymous()).to_string(), settings = ?settings, canister_id = %canister_id.to_string()))]
877    pub async fn create_canister_with_id(
878        &self,
879        sender: Option<Principal>,
880        settings: Option<CanisterSettings>,
881        canister_id: CanisterId,
882    ) -> Result<CanisterId, ClientError> {
883        self.create_canister_with_params(
884            sender,
885            CreateCanisterParams {
886                settings,
887                placement: Some(CreateCanisterPlacement::CanisterId(canister_id)),
888                ..Default::default()
889            },
890        )
891        .await
892    }
893
894    /// Create a canister on a specific subnet with optional custom settings.
895    /// The canister is created with 100T cycles.
896    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, sender = %sender.unwrap_or(Principal::anonymous()).to_string(), settings = ?settings, subnet_id = %subnet_id.to_string()))]
897    pub async fn create_canister_on_subnet(
898        &self,
899        sender: Option<Principal>,
900        settings: Option<CanisterSettings>,
901        subnet_id: SubnetId,
902    ) -> Result<CanisterId, ClientError> {
903        self.create_canister_with_params(
904            sender,
905            CreateCanisterParams {
906                settings,
907                placement: Some(CreateCanisterPlacement::SubnetId(subnet_id)),
908                ..Default::default()
909            },
910        )
911        .await
912    }
913
914    /// Create a canister with optional cycles, settings, and placement.
915    /// The placement specifies either a target subnet or a specific canister ID.
916    /// Defaults to 100T cycles if `params.cycles` is `None`.
917    /// Returns an error if the specified canister ID is already in use.
918    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
919    pub async fn create_canister_with_params(
920        &self,
921        sender: Option<Principal>,
922        params: CreateCanisterParams,
923    ) -> Result<CanisterId, ClientError> {
924        let cycles = params.cycles.unwrap_or(100_000_000_000_000);
925        let (effective_principal, specified_id) = match params.placement {
926            None => (RawEffectivePrincipal::None, None),
927            Some(CreateCanisterPlacement::SubnetId(subnet_id)) => (
928                RawEffectivePrincipal::SubnetId(subnet_id.as_slice().to_vec()),
929                None,
930            ),
931            Some(CreateCanisterPlacement::CanisterId(canister_id)) => (
932                RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
933                Some(canister_id),
934            ),
935        };
936        Ok(call_candid_as::<_, (CanisterIdRecord,)>(
937            self,
938            Principal::management_canister(),
939            effective_principal,
940            sender.unwrap_or(Principal::anonymous()),
941            "provisional_create_canister_with_cycles",
942            (ProvisionalCreateCanisterWithCyclesArgs {
943                settings: params.settings,
944                amount: Some(cycles.into()),
945                specified_id,
946                sender_canister_version: None,
947            },),
948        )
949        .await?
950        .0
951        .canister_id)
952    }
953
954    /// Upload a WASM chunk to the WASM chunk store of a canister.
955    /// Returns the WASM chunk hash.
956    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
957    pub async fn upload_chunk(
958        &self,
959        canister_id: CanisterId,
960        sender: Option<Principal>,
961        chunk: Vec<u8>,
962    ) -> Result<Vec<u8>, ClientError> {
963        Ok(call_candid_as::<_, (UploadChunkResult,)>(
964            self,
965            Principal::management_canister(),
966            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
967            sender.unwrap_or(Principal::anonymous()),
968            "upload_chunk",
969            (UploadChunkArgs { canister_id, chunk },),
970        )
971        .await?
972        .0
973        .hash)
974    }
975
976    /// List WASM chunk hashes in the WASM chunk store of a canister.
977    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
978    pub async fn stored_chunks(
979        &self,
980        canister_id: CanisterId,
981        sender: Option<Principal>,
982    ) -> Result<Vec<Vec<u8>>, ClientError> {
983        Ok(call_candid_as::<_, (StoredChunksResult,)>(
984            self,
985            Principal::management_canister(),
986            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
987            sender.unwrap_or(Principal::anonymous()),
988            "stored_chunks",
989            (CanisterIdRecord { canister_id },),
990        )
991        .await?
992        .0
993        .into_iter()
994        .map(|chunk| chunk.hash)
995        .collect())
996    }
997
998    /// Clear the WASM chunk store of a canister.
999    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1000    pub async fn clear_chunk_store(
1001        &self,
1002        canister_id: CanisterId,
1003        sender: Option<Principal>,
1004    ) -> Result<(), ClientError> {
1005        call_candid_as(
1006            self,
1007            Principal::management_canister(),
1008            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1009            sender.unwrap_or(Principal::anonymous()),
1010            "clear_chunk_store",
1011            (CanisterIdRecord { canister_id },),
1012        )
1013        .await
1014    }
1015
1016    /// Install a WASM module assembled from chunks on an existing canister.
1017    #[instrument(skip(self, mode, chunk_hashes_list, wasm_module_hash, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string(), store_canister_id = %store_canister_id.to_string(), arg_len = %arg.len()))]
1018    pub async fn install_chunked_canister(
1019        &self,
1020        canister_id: CanisterId,
1021        sender: Option<Principal>,
1022        mode: CanisterInstallMode,
1023        store_canister_id: CanisterId,
1024        chunk_hashes_list: Vec<Vec<u8>>,
1025        wasm_module_hash: Vec<u8>,
1026        arg: Vec<u8>,
1027    ) -> Result<(), ClientError> {
1028        call_candid_as(
1029            self,
1030            Principal::management_canister(),
1031            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1032            sender.unwrap_or(Principal::anonymous()),
1033            "install_chunked_code",
1034            (InstallChunkedCodeArgs {
1035                mode,
1036                target_canister: canister_id,
1037                store_canister: Some(store_canister_id),
1038                chunk_hashes_list: chunk_hashes_list
1039                    .into_iter()
1040                    .map(|hash| ChunkHash { hash })
1041                    .collect(),
1042                wasm_module_hash,
1043                arg,
1044                sender_canister_version: None,
1045            },),
1046        )
1047        .await
1048    }
1049
1050    async fn install_canister_helper(
1051        &self,
1052        mode: CanisterInstallMode,
1053        canister_id: CanisterId,
1054        wasm_module: Vec<u8>,
1055        arg: Vec<u8>,
1056        sender: Option<Principal>,
1057    ) -> Result<(), ClientError> {
1058        if wasm_module.len() + arg.len() < INSTALL_CHUNKED_CODE_THRESHOLD {
1059            call_candid_as::<(InstallCodeArgs,), ()>(
1060                self,
1061                Principal::management_canister(),
1062                RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1063                sender.unwrap_or(Principal::anonymous()),
1064                "install_code",
1065                (InstallCodeArgs {
1066                    mode,
1067                    canister_id,
1068                    wasm_module,
1069                    arg,
1070                    sender_canister_version: None,
1071                },),
1072            )
1073            .await?
1074        } else {
1075            self.clear_chunk_store(canister_id, sender).await?;
1076            let chunks: Vec<_> = wasm_module.chunks(INSTALL_CODE_CHUNK_SIZE).collect();
1077            let mut hashes = vec![];
1078            for chunk in chunks {
1079                let hash = self
1080                    .upload_chunk(canister_id, sender, chunk.to_vec())
1081                    .await?;
1082                hashes.push(hash);
1083            }
1084            let mut hasher = Sha256::new();
1085            hasher.update(wasm_module);
1086            let wasm_module_hash = hasher.finalize().to_vec();
1087            self.install_chunked_canister(
1088                canister_id,
1089                sender,
1090                mode,
1091                canister_id,
1092                hashes,
1093                wasm_module_hash,
1094                arg,
1095            )
1096            .await?
1097        };
1098        Ok(())
1099    }
1100
1101    /// Install a WASM module on an existing canister.
1102    #[instrument(skip(self, wasm_module, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1103    pub async fn install_canister(
1104        &self,
1105        canister_id: CanisterId,
1106        wasm_module: Vec<u8>,
1107        arg: Vec<u8>,
1108        sender: Option<Principal>,
1109    ) -> Result<(), ClientError> {
1110        self.install_canister_helper(
1111            CanisterInstallMode::Install,
1112            canister_id,
1113            wasm_module,
1114            arg,
1115            sender,
1116        )
1117        .await
1118    }
1119
1120    /// Upgrade a canister with a new WASM module.
1121    #[instrument(skip(self, wasm_module, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1122    pub async fn upgrade_canister(
1123        &self,
1124        canister_id: CanisterId,
1125        wasm_module: Vec<u8>,
1126        arg: Vec<u8>,
1127        sender: Option<Principal>,
1128    ) -> Result<(), ClientError> {
1129        self.install_canister_helper(
1130            CanisterInstallMode::Upgrade(None),
1131            canister_id,
1132            wasm_module,
1133            arg,
1134            sender,
1135        )
1136        .await
1137    }
1138
1139    /// Reinstall a canister WASM module.
1140    #[instrument(skip(self, wasm_module, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1141    pub async fn reinstall_canister(
1142        &self,
1143        canister_id: CanisterId,
1144        wasm_module: Vec<u8>,
1145        arg: Vec<u8>,
1146        sender: Option<Principal>,
1147    ) -> Result<(), ClientError> {
1148        self.install_canister_helper(
1149            CanisterInstallMode::Reinstall,
1150            canister_id,
1151            wasm_module,
1152            arg,
1153            sender,
1154        )
1155        .await
1156    }
1157
1158    /// Uninstall a canister.
1159    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1160    pub async fn uninstall_canister(
1161        &self,
1162        canister_id: CanisterId,
1163        sender: Option<Principal>,
1164    ) -> Result<(), ClientError> {
1165        call_candid_as::<(CanisterIdRecord,), ()>(
1166            self,
1167            Principal::management_canister(),
1168            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1169            sender.unwrap_or(Principal::anonymous()),
1170            "uninstall_code",
1171            (CanisterIdRecord { canister_id },),
1172        )
1173        .await
1174    }
1175
1176    /// Take canister snapshot.
1177    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1178    pub async fn take_canister_snapshot(
1179        &self,
1180        canister_id: CanisterId,
1181        sender: Option<Principal>,
1182        replace_snapshot: Option<Vec<u8>>,
1183    ) -> Result<Snapshot, ClientError> {
1184        Ok(call_candid_as::<_, (Snapshot,)>(
1185            self,
1186            Principal::management_canister(),
1187            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1188            sender.unwrap_or(Principal::anonymous()),
1189            "take_canister_snapshot",
1190            (TakeCanisterSnapshotArgs {
1191                canister_id,
1192                replace_snapshot,
1193                uninstall_code: None,
1194                sender_canister_version: None,
1195            },),
1196        )
1197        .await?
1198        .0)
1199    }
1200
1201    /// Load canister snapshot.
1202    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1203    pub async fn load_canister_snapshot(
1204        &self,
1205        canister_id: CanisterId,
1206        sender: Option<Principal>,
1207        snapshot_id: Vec<u8>,
1208    ) -> Result<(), ClientError> {
1209        call_candid_as(
1210            self,
1211            Principal::management_canister(),
1212            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1213            sender.unwrap_or(Principal::anonymous()),
1214            "load_canister_snapshot",
1215            (LoadCanisterSnapshotArgs {
1216                canister_id,
1217                snapshot_id,
1218                sender_canister_version: None,
1219            },),
1220        )
1221        .await
1222    }
1223
1224    /// List canister snapshots.
1225    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1226    pub async fn list_canister_snapshots(
1227        &self,
1228        canister_id: CanisterId,
1229        sender: Option<Principal>,
1230    ) -> Result<Vec<Snapshot>, ClientError> {
1231        Ok(call_candid_as::<_, (Vec<Snapshot>,)>(
1232            self,
1233            Principal::management_canister(),
1234            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1235            sender.unwrap_or(Principal::anonymous()),
1236            "list_canister_snapshots",
1237            (CanisterIdRecord { canister_id },),
1238        )
1239        .await?
1240        .0)
1241    }
1242
1243    /// Delete canister snapshot.
1244    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1245    pub async fn delete_canister_snapshot(
1246        &self,
1247        canister_id: CanisterId,
1248        sender: Option<Principal>,
1249        snapshot_id: Vec<u8>,
1250    ) -> Result<(), ClientError> {
1251        call_candid_as(
1252            self,
1253            Principal::management_canister(),
1254            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1255            sender.unwrap_or(Principal::anonymous()),
1256            "delete_canister_snapshot",
1257            (DeleteCanisterSnapshotArgs {
1258                canister_id,
1259                snapshot_id,
1260            },),
1261        )
1262        .await
1263    }
1264
1265    /// Update canister settings.
1266    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1267    pub async fn update_canister_settings(
1268        &self,
1269        canister_id: CanisterId,
1270        sender: Option<Principal>,
1271        settings: CanisterSettings,
1272    ) -> Result<(), ClientError> {
1273        call_candid_as::<_, ()>(
1274            self,
1275            Principal::management_canister(),
1276            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1277            sender.unwrap_or(Principal::anonymous()),
1278            "update_settings",
1279            (UpdateSettingsArgs {
1280                canister_id,
1281                settings,
1282                sender_canister_version: None,
1283            },),
1284        )
1285        .await
1286    }
1287
1288    /// Set canister's controllers.
1289    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1290    pub async fn set_controllers(
1291        &self,
1292        canister_id: CanisterId,
1293        sender: Option<Principal>,
1294        new_controllers: Vec<Principal>,
1295    ) -> Result<(), ClientError> {
1296        let settings = CanisterSettings {
1297            controllers: Some(new_controllers),
1298            ..CanisterSettings::default()
1299        };
1300        call_candid_as::<(UpdateSettingsArgs,), ()>(
1301            self,
1302            Principal::management_canister(),
1303            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1304            sender.unwrap_or(Principal::anonymous()),
1305            "update_settings",
1306            (UpdateSettingsArgs {
1307                canister_id,
1308                settings,
1309                sender_canister_version: None,
1310            },),
1311        )
1312        .await
1313    }
1314
1315    /// Start a canister.
1316    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1317    pub async fn start_canister(
1318        &self,
1319        canister_id: CanisterId,
1320        sender: Option<Principal>,
1321    ) -> Result<(), ClientError> {
1322        call_candid_as::<(CanisterIdRecord,), ()>(
1323            self,
1324            Principal::management_canister(),
1325            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1326            sender.unwrap_or(Principal::anonymous()),
1327            "start_canister",
1328            (CanisterIdRecord { canister_id },),
1329        )
1330        .await
1331    }
1332
1333    /// Stop a canister.
1334    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1335    pub async fn stop_canister(
1336        &self,
1337        canister_id: CanisterId,
1338        sender: Option<Principal>,
1339    ) -> Result<(), ClientError> {
1340        call_candid_as::<(CanisterIdRecord,), ()>(
1341            self,
1342            Principal::management_canister(),
1343            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1344            sender.unwrap_or(Principal::anonymous()),
1345            "stop_canister",
1346            (CanisterIdRecord { canister_id },),
1347        )
1348        .await
1349    }
1350
1351    /// Delete a canister.
1352    #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
1353    pub async fn delete_canister(
1354        &self,
1355        canister_id: CanisterId,
1356        sender: Option<Principal>,
1357    ) -> Result<(), ClientError> {
1358        call_candid_as::<(CanisterIdRecord,), ()>(
1359            self,
1360            Principal::management_canister(),
1361            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
1362            sender.unwrap_or(Principal::anonymous()),
1363            "delete_canister",
1364            (CanisterIdRecord { canister_id },),
1365        )
1366        .await
1367    }
1368
1369    /// Checks whether the provided canister exists.
1370    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
1371    pub async fn canister_exists(&self, canister_id: CanisterId) -> Result<bool, ClientError> {
1372        Ok(self.get_subnet(canister_id).await?.is_some())
1373    }
1374
1375    /// Deletes a subnet. Returns a server error if the subnet does not exist or is a named subnet.
1376    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, subnet_id = %subnet_id.to_string()))]
1377    pub async fn delete_subnet(&self, subnet_id: SubnetId) -> Result<(), ClientError> {
1378        let endpoint = "update/delete_subnet";
1379        self.post::<(), RawSubnetId>(
1380            endpoint,
1381            RawSubnetId {
1382                subnet_id: subnet_id.as_slice().to_vec(),
1383            },
1384        )
1385        .await?;
1386
1387        Ok(())
1388    }
1389
1390    /// Returns the subnet ID of the canister if the canister exists.
1391    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string()))]
1392    pub async fn get_subnet(
1393        &self,
1394        canister_id: CanisterId,
1395    ) -> Result<Option<SubnetId>, ClientError> {
1396        let endpoint = "read/get_subnet";
1397        let result: Option<RawSubnetId> = self
1398            .post(
1399                endpoint,
1400                RawCanisterId {
1401                    canister_id: canister_id.as_slice().to_vec(),
1402                },
1403            )
1404            .await?;
1405        Ok(result.map(|RawSubnetId { subnet_id }| SubnetId::from_slice(&subnet_id)))
1406    }
1407
1408    /// Returns subnet metrics for a given subnet.
1409    #[instrument(ret, skip(self), fields(instance_id=self.instance_id, subnet_id = %subnet_id.to_string()))]
1410    pub async fn get_subnet_metrics(
1411        &self,
1412        subnet_id: Principal,
1413    ) -> Result<Option<SubnetMetrics>, ClientError> {
1414        let path = vec![
1415            "subnet".into(),
1416            Label::from_bytes(subnet_id.as_slice()),
1417            "metrics".into(),
1418        ];
1419        let paths = vec![path.clone()];
1420        let content = ReadState {
1421            ingress_expiry: self.get_time().await?.as_nanos_since_unix_epoch() + 240_000_000_000,
1422            sender: Principal::anonymous(),
1423            paths,
1424        };
1425        let envelope = Envelope {
1426            content: std::borrow::Cow::Borrowed(&content),
1427            sender_pubkey: None,
1428            sender_sig: None,
1429            sender_delegation: None,
1430        };
1431
1432        let serialized_bytes = encode_envelope(&envelope, usize::MAX)?;
1433
1434        let endpoint = format!("api/v2/subnet/{}/read_state", subnet_id.to_text());
1435        let url = self.instance_url().join(&endpoint).unwrap();
1436        let mut exchange = Exchange::new(&url, self.max_request_time_ms);
1437        let (status, bytes) = exchange
1438            .bytes(
1439                &self.reqwest_client,
1440                self.reqwest_client
1441                    .post(url.clone())
1442                    .header(reqwest::header::CONTENT_TYPE, "application/cbor")
1443                    .body(serialized_bytes),
1444            )
1445            .await?;
1446        if !status.is_success() {
1447            return Err(exchange.server_error(status, &bytes));
1448        }
1449        let read_state_response: ReadStateResponse =
1450            decode_read_state_response(&bytes, bytes.len())?;
1451        let cert: Certificate = decode_certificate(
1452            &read_state_response.certificate,
1453            read_state_response.certificate.len(),
1454        )?;
1455
1456        let metrics = match cert.tree.lookup_path(path) {
1457            LookupResult::Found(value) => value,
1458            _ => return Ok(None),
1459        };
1460        decode_subnet_metrics(metrics, metrics.len()).map_err(ClientError::from)
1461    }
1462
1463    /// This (asynchronous) drop function must be called to drop the PocketIc instance.
1464    /// It must be called manually as Rust doesn't support asynchronous drop.
1465    pub async fn drop(mut self) -> Result<(), ClientError> {
1466        self.do_drop().await?;
1467
1468        Ok(())
1469    }
1470
1471    pub(crate) async fn do_drop(&mut self) -> Result<(), ClientError> {
1472        if self.owns_instance {
1473            let url = self.instance_url();
1474            let mut exchange = Exchange::new(&url, self.max_request_time_ms);
1475            let (status, bytes) = exchange
1476                .bytes(
1477                    &self.reqwest_client,
1478                    self.reqwest_client.delete(url.clone()),
1479                )
1480                .await?;
1481            if !status.is_success() {
1482                return Err(exchange.server_error(status, &bytes));
1483            }
1484            self.owns_instance = false;
1485            Ok(())
1486        } else {
1487            self.stop_http_gateway().await
1488        }
1489    }
1490
1491    async fn get<T: DeserializeOwned>(&self, endpoint: &str) -> Result<T, ClientError> {
1492        self.request(HttpMethod::Get, endpoint, ()).await
1493    }
1494
1495    async fn post<T: DeserializeOwned, B: Serialize>(
1496        &self,
1497        endpoint: &str,
1498        body: B,
1499    ) -> Result<T, ClientError> {
1500        self.request(HttpMethod::Post, endpoint, body).await
1501    }
1502
1503    async fn request<T: DeserializeOwned, B: Serialize>(
1504        &self,
1505        method: HttpMethod,
1506        endpoint: &str,
1507        body: B,
1508    ) -> Result<T, ClientError> {
1509        self.request_limited(method, endpoint, body, None).await
1510    }
1511
1512    async fn request_limited<T: DeserializeOwned, B: Serialize>(
1513        &self,
1514        method: HttpMethod,
1515        endpoint: &str,
1516        body: B,
1517        response_limit: Option<usize>,
1518    ) -> Result<T, ClientError> {
1519        let url = self.instance_url().join(endpoint).unwrap();
1520        let mut exchange = Exchange::new(&url, self.max_request_time_ms);
1521        exchange.response_limit = response_limit;
1522        loop {
1523            let builder = if let Some(operation) = &exchange.context.operation {
1524                self.reqwest_client.get(
1525                    self.server_url
1526                        .join(&format!(
1527                            "/read_graph/{}/{}",
1528                            operation.state_label, operation.op_id
1529                        ))
1530                        .unwrap(),
1531                )
1532            } else {
1533                match method {
1534                    HttpMethod::Get => self.reqwest_client.get(url.clone()),
1535                    HttpMethod::Post => self.reqwest_client.post(url.clone()).json(&body),
1536                }
1537            };
1538            let (status, bytes) = exchange.bytes(&self.reqwest_client, builder).await?;
1539            match status {
1540                StatusCode::OK => {
1541                    let value = exchange.decode(&bytes)?;
1542                    if let Some(operation) = &exchange.context.operation {
1543                        // The result is known; pruning is best-effort housekeeping, never a resubmission.
1544                        let request = self.reqwest_client.delete(
1545                            self.server_url
1546                                .join(&format!(
1547                                    "/prune_graph/{}/{}",
1548                                    operation.state_label, operation.op_id
1549                                ))
1550                                .unwrap(),
1551                        );
1552                        tokio::spawn(async move {
1553                            let _ = request.send().await;
1554                        });
1555                    }
1556                    return Ok(value);
1557                }
1558                StatusCode::ACCEPTED | StatusCode::CONFLICT => {
1559                    let receipt: crate::pocket_ic::common::rest::StartedOrBusyResponse =
1560                        exchange.decode(&bytes)?;
1561                    if exchange.context.operation.is_none() {
1562                        if status == StatusCode::ACCEPTED {
1563                            exchange.started(Operation {
1564                                state_label: receipt.state_label,
1565                                op_id: receipt.op_id,
1566                            });
1567                        } else {
1568                            // Busy explicitly says this submission was not accepted.
1569                            exchange.context.may_have_executed = false;
1570                        }
1571                    }
1572                    // Once accepted, always poll our original receipt, even if an
1573                    // unexpected response refers to a different computation.
1574                }
1575                StatusCode::NOT_FOUND if exchange.context.operation.is_some() => {}
1576                _ => return Err(exchange.server_error(status, &bytes)),
1577            }
1578            exchange.pause().await?;
1579        }
1580    }
1581
1582    /// Get the pending canister HTTP outcalls.
1583    /// Note that an additional `PocketIc::tick` is necessary after a canister
1584    /// executes a message making a canister HTTP outcall for the HTTP outcall
1585    /// to be retrievable here.
1586    /// Note that, unless a PocketIC instance is in auto progress mode,
1587    /// a response to the pending canister HTTP outcalls
1588    /// must be produced by the test driver and passed on to the PocketIC instace
1589    /// using `PocketIc::mock_canister_http_response`. Flexible outcall mocking is
1590    /// not supported by this client.
1591    /// In auto progress mode, the PocketIC server produces a response for every
1592    /// pending canister HTTP outcall by actually making an HTTP request
1593    /// to the specified URL.
1594    #[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
1595    pub async fn get_canister_http(&self) -> Result<Vec<CanisterHttpRequest>, ClientError> {
1596        let endpoint = "read/get_canister_http";
1597        let res: Vec<RawCanisterHttpRequest> = self.get(endpoint).await?;
1598        Ok(res.into_iter().map(|r| r.into()).collect())
1599    }
1600
1601    /// Mock a response to a pending canister HTTP outcall: the same response for
1602    /// every node of the subnet, or one response per node if
1603    /// `MockCanisterHttpResponse::additional_responses` is non-empty.
1604    #[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
1605    pub async fn mock_canister_http_response(
1606        &self,
1607        mock_canister_http_response: MockCanisterHttpResponse,
1608    ) -> Result<(), ClientError> {
1609        let endpoint = "update/mock_canister_http";
1610        let raw_mock_canister_http_response: RawMockCanisterHttpResponse =
1611            mock_canister_http_response.into();
1612        self.post(endpoint, raw_mock_canister_http_response).await
1613    }
1614
1615    /// Download a canister snapshot to a given snapshot directory.
1616    /// The sender must be a controller of the canister.
1617    /// The snapshot directory must be empty if it exists.
1618    #[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
1619    pub async fn canister_snapshot_download(
1620        &self,
1621        canister_id: CanisterId,
1622        sender: Principal,
1623        snapshot_id: Vec<u8>,
1624        snapshot_dir: PathBuf,
1625    ) -> Result<(), ClientError> {
1626        let endpoint = "update/canister_snapshot_download";
1627        #[cfg(not(windows))]
1628        let snapshot_dir = snapshot_dir;
1629        #[cfg(windows)]
1630        let snapshot_dir = wsl_path(&snapshot_dir, "snapshot directory").into();
1631        let raw_canister_snapshot_download = RawCanisterSnapshotDownload {
1632            sender: sender.into(),
1633            canister_id: canister_id.into(),
1634            snapshot_id,
1635            snapshot_dir,
1636        };
1637        self.post(endpoint, raw_canister_snapshot_download).await
1638    }
1639
1640    /// Upload a canister snapshot from a given snapshot directory.
1641    /// The sender must be a controller of the canister.
1642    /// Returns the snapshot ID of the uploaded snapshot.
1643    #[instrument(ret, skip(self), fields(instance_id=self.instance_id))]
1644    pub async fn canister_snapshot_upload(
1645        &self,
1646        canister_id: CanisterId,
1647        sender: Principal,
1648        replace_snapshot: Option<Vec<u8>>,
1649        snapshot_dir: PathBuf,
1650    ) -> Result<Vec<u8>, ClientError> {
1651        let endpoint = "update/canister_snapshot_upload";
1652        let replace_snapshot =
1653            replace_snapshot.map(|snapshot_id| RawCanisterSnapshotId { snapshot_id });
1654        #[cfg(not(windows))]
1655        let snapshot_dir = snapshot_dir;
1656        #[cfg(windows)]
1657        let snapshot_dir = wsl_path(&snapshot_dir, "snapshot directory").into();
1658        let raw_canister_snapshot_upload = RawCanisterSnapshotUpload {
1659            sender: sender.into(),
1660            canister_id: canister_id.into(),
1661            replace_snapshot,
1662            snapshot_dir,
1663        };
1664        Ok(self
1665            .post::<RawCanisterSnapshotId, _>(endpoint, raw_canister_snapshot_upload)
1666            .await?
1667            .snapshot_id)
1668    }
1669}
1670
1671/// Call a canister candid method, authenticated. The sender can be impersonated (i.e., the
1672/// signature is not verified).
1673/// PocketIC executes update calls synchronously, so there is no need to poll for the result.
1674pub async fn call_candid_as<Input, Output>(
1675    env: &PocketIc,
1676    canister_id: CanisterId,
1677    effective_principal: RawEffectivePrincipal,
1678    sender: Principal,
1679    method: &str,
1680    input: Input,
1681) -> Result<Output, ClientError>
1682where
1683    Input: ArgumentEncoder,
1684    Output: for<'a> ArgumentDecoder<'a>,
1685{
1686    with_candid(input, |payload| async {
1687        env.update_call_with_effective_principal(
1688            canister_id,
1689            effective_principal,
1690            sender,
1691            method,
1692            payload,
1693        )
1694        .await
1695    })
1696    .await
1697}
1698
1699/// Call a canister candid method, anonymous.
1700/// PocketIC executes update calls synchronously, so there is no need to poll for the result.
1701pub async fn call_candid<Input, Output>(
1702    env: &PocketIc,
1703    canister_id: CanisterId,
1704    effective_principal: RawEffectivePrincipal,
1705    method: &str,
1706    input: Input,
1707) -> Result<Output, ClientError>
1708where
1709    Input: ArgumentEncoder,
1710    Output: for<'a> ArgumentDecoder<'a>,
1711{
1712    call_candid_as(
1713        env,
1714        canister_id,
1715        effective_principal,
1716        Principal::anonymous(),
1717        method,
1718        input,
1719    )
1720    .await
1721}
1722
1723/// Call a canister candid query method, anonymous.
1724pub async fn query_candid<Input, Output>(
1725    env: &PocketIc,
1726    canister_id: CanisterId,
1727    method: &str,
1728    input: Input,
1729) -> Result<Output, ClientError>
1730where
1731    Input: ArgumentEncoder,
1732    Output: for<'a> ArgumentDecoder<'a>,
1733{
1734    query_candid_as(env, canister_id, Principal::anonymous(), method, input).await
1735}
1736
1737/// Call a canister candid query method, authenticated. The sender can be impersonated (i.e., the
1738/// signature is not verified).
1739pub async fn query_candid_as<Input, Output>(
1740    env: &PocketIc,
1741    canister_id: CanisterId,
1742    sender: Principal,
1743    method: &str,
1744    input: Input,
1745) -> Result<Output, ClientError>
1746where
1747    Input: ArgumentEncoder,
1748    Output: for<'a> ArgumentDecoder<'a>,
1749{
1750    with_candid(input, |bytes| async {
1751        env.query_call(canister_id, sender, method, bytes).await
1752    })
1753    .await
1754}
1755
1756/// Call a canister candid update method, anonymous.
1757pub async fn update_candid<Input, Output>(
1758    env: &PocketIc,
1759    canister_id: CanisterId,
1760    method: &str,
1761    input: Input,
1762) -> Result<Output, ClientError>
1763where
1764    Input: ArgumentEncoder,
1765    Output: for<'a> ArgumentDecoder<'a>,
1766{
1767    update_candid_as(env, canister_id, Principal::anonymous(), method, input).await
1768}
1769
1770/// Call a canister candid update method, authenticated. The sender can be impersonated (i.e., the
1771/// signature is not verified).
1772pub async fn update_candid_as<Input, Output>(
1773    env: &PocketIc,
1774    canister_id: CanisterId,
1775    sender: Principal,
1776    method: &str,
1777    input: Input,
1778) -> Result<Output, ClientError>
1779where
1780    Input: ArgumentEncoder,
1781    Output: for<'a> ArgumentDecoder<'a>,
1782{
1783    with_candid(input, |bytes| async {
1784        env.update_call(canister_id, sender, method, bytes).await
1785    })
1786    .await
1787}
1788
1789/// A helper function that we use to implement both [`call_candid`] and
1790/// [`query_candid`].
1791pub async fn with_candid<Input, Output, Fut>(
1792    input: Input,
1793    f: impl FnOnce(Vec<u8>) -> Fut,
1794) -> Result<Output, ClientError>
1795where
1796    Input: ArgumentEncoder,
1797    Output: for<'a> ArgumentDecoder<'a>,
1798    Fut: Future<Output = Result<Vec<u8>, ClientError>>,
1799{
1800    let input = encode_args(input)?;
1801    let output = f(input).await?;
1802    Ok(decode_args(&output)?)
1803}
1804
1805fn setup_tracing(pid: u32) -> Option<WorkerGuard> {
1806    use tracing_subscriber::prelude::*;
1807    match std::env::var(LOG_DIR_PATH_ENV_NAME).map(std::path::PathBuf::from) {
1808        Ok(p) => {
1809            std::fs::create_dir_all(&p).expect("Could not create directory");
1810
1811            let file_name = format!("pocket_ic_client_{pid}.log");
1812            let appender = tracing_appender::rolling::never(&p, file_name);
1813            let (non_blocking_appender, guard) = tracing_appender::non_blocking(appender);
1814            let log_dir_filter: EnvFilter =
1815                tracing_subscriber::EnvFilter::try_from_env(LOG_DIR_LEVELS_ENV_NAME)
1816                    .unwrap_or_else(|_| "trace".parse().unwrap());
1817
1818            let layers = vec![
1819                tracing_subscriber::fmt::layer()
1820                    .with_writer(non_blocking_appender)
1821                    // disable color escape codes in files
1822                    .with_ansi(false)
1823                    .with_filter(log_dir_filter)
1824                    .boxed(),
1825            ];
1826            let _ = tracing_subscriber::registry().with(layers).try_init();
1827            Some(guard)
1828        }
1829        _ => None,
1830    }
1831}
1832
1833/// Retrieves a default effective canister id for canister creation on a PocketIC instance
1834/// characterized by:
1835///  - a PocketIC instance URL of the form `http://<ip>:<port>/instances/<instance_id>`;
1836///  - a PocketIC HTTP gateway URL of the form `http://<ip>:port` for a PocketIC instance.
1837///
1838/// Returns an error if the PocketIC instance topology could not be fetched or parsed, e.g.,
1839/// because the given URL points to a replica (i.e., does not meet any of the above two properties).
1840pub async fn get_default_effective_canister_id(
1841    pocket_ic_url: String,
1842) -> Result<Principal, ClientError> {
1843    let url = Url::parse(&format!(
1844        "{}/_/topology",
1845        pocket_ic_url.trim_end_matches('/')
1846    ))
1847    .map_err(|e| ClientError::Protocol(e.to_string()))?;
1848    let client = reqwest::Client::new();
1849    let mut exchange = Exchange::new(&url, Some(crate::pocket_ic::DEFAULT_MAX_REQUEST_TIME_MS));
1850    loop {
1851        let (status, bytes) = exchange.bytes(&client, client.get(url.clone())).await?;
1852        if status == StatusCode::CONFLICT {
1853            exchange.pause().await?;
1854            continue;
1855        }
1856        if !status.is_success() {
1857            return Err(exchange.server_error(status, &bytes));
1858        }
1859        let topology: Topology = exchange.decode(&bytes)?;
1860        return Principal::try_from_slice(&topology.default_effective_canister_id.canister_id)
1861            .map_err(|e| ClientError::Protocol(e.to_string()));
1862    }
1863}