1#![forbid(unsafe_code)]
4#![forbid(missing_docs)]
5
6use async_trait::async_trait;
7use candid::{Decode, Encode};
8use ic_http_types::{HttpRequest, HttpResponse};
9#[cfg(feature = "pocket_ic")]
10pub use pocket_ic_query_call::{PocketIcAsyncHttpQuery, PocketIcHttpQuery};
11use regex::Regex;
12use std::fmt;
13use std::fmt::Debug;
14
15#[cfg_attr(feature = "pocket_ic", doc = "```rust")]
20#[cfg_attr(not(feature = "pocket_ic"), doc = "```ignore")]
21pub struct MetricsAssert<T> {
62 actual: T,
63 metrics: Vec<String>,
64}
65
66impl<T> MetricsAssert<T> {
67 pub fn from_http_query<E>(actual: T) -> Self
70 where
71 T: CanisterHttpQuery<E>,
72 E: Debug,
73 {
74 let metrics =
75 decode_metrics_response_or_unwrap(actual.http_query(encoded_metrics_request()));
76 Self { actual, metrics }
77 }
78
79 pub async fn from_async_http_query<E>(actual: T) -> Self
82 where
83 T: AsyncCanisterHttpQuery<E>,
84 E: Debug,
85 {
86 let metrics =
87 decode_metrics_response_or_unwrap(actual.http_query(encoded_metrics_request()).await);
88 Self { actual, metrics }
89 }
90
91 pub fn into(self) -> T {
93 self.actual
94 }
95
96 pub fn assert_contains_metric_matching<P: AsRef<str> + fmt::Display>(self, pattern: P) -> Self {
98 assert!(
99 !self.find_metrics_matching(pattern.as_ref()).is_empty(),
100 "Expected to find metric matching '{}', but none matched in:\n{:?}",
101 pattern,
102 self.metrics
103 );
104 self
105 }
106
107 pub fn assert_does_not_contain_metric_matching(self, pattern: &str) -> Self {
109 let matches = self.find_metrics_matching(pattern);
110 assert!(
111 matches.is_empty(),
112 "Expected not to find any metric matching '{pattern}', but found the following matches:\n{matches:?}"
113 );
114 self
115 }
116
117 pub fn find_metrics_matching(&self, pattern: &str) -> Vec<String> {
119 let regex = Regex::new(pattern).unwrap_or_else(|_| panic!("Invalid regex: {pattern}"));
120 self.metrics
121 .iter()
122 .filter(|line| regex.is_match(line))
123 .cloned()
124 .collect()
125 }
126}
127
128fn encoded_metrics_request() -> Vec<u8> {
129 let request = HttpRequest {
130 method: "GET".to_string(),
131 url: "/metrics".to_string(),
132 headers: Default::default(),
133 body: Default::default(),
134 };
135 Encode!(&request).expect("failed to encode HTTP request")
136}
137
138fn decode_metrics_response_or_unwrap<E: Debug>(response: Result<Vec<u8>, E>) -> Vec<String> {
139 let response = Decode!(&response.expect("failed to retrieve metrics"), HttpResponse)
140 .expect("failed to decode HTTP response");
141 assert_eq!(response.status_code, 200_u16);
142 String::from_utf8_lossy(response.body.as_slice())
143 .trim()
144 .split('\n')
145 .map(|line| line.to_string())
146 .collect()
147}
148
149pub trait CanisterHttpQuery<E: Debug> {
151 fn http_query(&self, request: Vec<u8>) -> Result<Vec<u8>, E>;
153}
154
155#[async_trait]
157pub trait AsyncCanisterHttpQuery<E: Debug> {
158 async fn http_query(&self, request: Vec<u8>) -> Result<Vec<u8>, E>;
160}
161
162#[cfg(feature = "pocket_ic")]
163mod pocket_ic_query_call {
164 use super::*;
165 use candid::Principal;
166 use ic_management_canister_types::CanisterId;
167 use pocket_ic::{PocketIc, RejectResponse, nonblocking};
168
169 pub trait PocketIcHttpQuery {
172 fn get_pocket_ic(&self) -> &PocketIc;
174
175 fn get_canister_id(&self) -> CanisterId;
177 }
178
179 impl<T: PocketIcHttpQuery> CanisterHttpQuery<RejectResponse> for T {
180 fn http_query(&self, request: Vec<u8>) -> Result<Vec<u8>, RejectResponse> {
181 self.get_pocket_ic().query_call(
182 self.get_canister_id(),
183 Principal::anonymous(),
184 "http_request",
185 request,
186 )
187 }
188 }
189
190 pub trait PocketIcAsyncHttpQuery {
193 fn get_pocket_ic(&self) -> &nonblocking::PocketIc;
196
197 fn get_canister_id(&self) -> CanisterId;
199 }
200
201 #[async_trait]
202 impl<T: PocketIcAsyncHttpQuery + Send + Sync> AsyncCanisterHttpQuery<RejectResponse> for T {
203 async fn http_query(&self, request: Vec<u8>) -> Result<Vec<u8>, RejectResponse> {
204 self.get_pocket_ic()
205 .query_call(
206 self.get_canister_id(),
207 Principal::anonymous(),
208 "http_request",
209 request,
210 )
211 .await
212 }
213 }
214}