pdk_unit/tester/
io.rs

1// Copyright (c) 2026, Salesforce, Inc.,
2// All rights reserved.
3// For full license text, see the LICENSE.txt file
4
5use classy::stream::PropertyAccessor;
6use pdk_core::policy_context::authentication::{
7    Authentication, AuthenticationData, AuthenticationHandler,
8};
9use pdk_core::policy_context::policy_violation::{PolicyViolation, PolicyViolations};
10use proxy_wasm_stub::types::Bytes;
11use serde::{Deserialize, Deserializer, Serialize, Serializer};
12use std::cell::RefCell;
13use std::collections::HashMap;
14use std::fmt::{Debug, Formatter};
15use std::rc::Rc;
16
17/// An HTTP request used in `pdk-unit` tests.
18///
19/// Construct one using the HTTP method constructors ([`get`](Self::get), [`post`](Self::post), etc.)
20/// or [`custom`](Self::custom) for non-standard methods, then chain `with_*` builder methods to
21/// populate headers, body, properties, and authentication.
22///
23/// Read-only accessors are available via the [`UnitHttpMessage`] trait.
24///
25/// # Example
26///
27/// ```ignore
28/// use pdk_unit::{UnitHttpRequest, UnitHttpMessage};
29///
30/// let req = UnitHttpRequest::get()
31///     .with_path("/api/users")
32///     .with_header("authorization", "Bearer token123")
33///     .with_body("hello");
34///
35/// assert_eq!(req.header("authorization"), Some("Bearer token123"));
36/// ```
37#[derive(Clone, Debug, PartialEq)]
38pub struct UnitHttpRequest {
39    pub(crate) inner: RequestResponse,
40}
41
42/// An HTTP response used in `pdk-unit` tests.
43///
44/// Construct one with [`new`](Self::new) providing the HTTP status code, then chain `with_*`
45/// builder methods to populate headers and body.
46///
47/// Read-only accessors are available via the [`UnitHttpMessage`] trait.
48///
49/// # Example
50///
51/// ```ignore
52/// use pdk_unit::{UnitHttpResponse, UnitHttpMessage};
53///
54/// let resp = UnitHttpResponse::new(200)
55///     .with_header("content-type", "application/json")
56///     .with_body(r#"{"ok":true}"#);
57///
58/// assert_eq!(resp.status_code(), 200);
59/// assert_eq!(resp.header("content-type"), Some("application/json"));
60/// ```
61#[derive(Clone, Debug, PartialEq)]
62pub struct UnitHttpResponse {
63    pub(crate) inner: RequestResponse,
64}
65
66macro_rules! http_method {
67    ($fn_name:ident, $method:expr) => {
68        #[doc = "Creates a `"]
69        #[doc = $method]
70        #[doc = "` request."]
71        pub fn $fn_name() -> Self {
72            Self::custom($method)
73        }
74    };
75}
76
77/// Common read-only accessors for HTTP messages in unit tests.
78///
79/// Implemented by both [`UnitHttpRequest`] and [`UnitHttpResponse`], allowing test helper
80/// functions to be written generically over either type.
81///
82/// # Example
83///
84/// ```ignore
85/// use pdk_unit::{UnitHttpMessage, UnitHttpRequest, UnitHttpResponse};
86///
87/// fn assert_ok<M: UnitHttpMessage>(msg: &M) {
88///     assert!(msg.body().len() > 0);
89/// }
90/// ```
91pub trait UnitHttpMessage {
92    /// Returns the value of the header with the given name, or `None` if not present.
93    fn header(&self, header: &str) -> Option<&str>;
94
95    /// Returns all headers as a list of `(name, value)` pairs.
96    fn headers(&self) -> &Vec<(String, String)>;
97
98    /// Returns the message body as a byte slice.
99    fn body(&self) -> &[u8];
100
101    /// Returns the value of a property identified by the given key path, or `None` if not set.
102    fn property<K: Into<String>>(&self, key: Vec<K>) -> Option<Bytes>;
103
104    /// Returns all properties as a map of key path to raw bytes.
105    fn properties(&self) -> HashMap<Vec<String>, Bytes>;
106
107    /// Returns the authentication data attached to this message, if any.
108    fn authentication(&self) -> Option<AuthenticationData>;
109
110    /// Returns the policy violation attached to this message, if any.
111    fn violation(&self) -> Option<PolicyViolation>;
112}
113
114macro_rules! impl_request_response_methods {
115    () => {
116        /// Sets a header, replacing any existing value for the same name.
117        pub fn with_header<K: Into<String>, V: Into<String>>(mut self, key: K, val: V) -> Self {
118            self.inner = self.inner.with_header(key, val);
119            self
120        }
121
122        /// Sets the message body.
123        pub fn with_body<B: Into<Vec<u8>>>(mut self, body: B) -> Self {
124            self.inner = self.inner.with_body(body);
125            self
126        }
127
128        /// Sets a single property identified by the given key path.
129        pub fn with_property<K: Into<String>, V: Into<Vec<u8>>>(
130            mut self,
131            key: Vec<K>,
132            value: V,
133        ) -> Self {
134            self.inner = self.inner.with_property(key, value);
135            self
136        }
137
138        /// Replaces all properties with the given map.
139        pub fn with_properties(mut self, properties: HashMap<Vec<String>, Bytes>) -> Self {
140            self.inner = self.inner.with_properties(properties);
141            self
142        }
143
144        /// Attaches authentication data to this message.
145        pub fn with_authentication_data(mut self, authentication: AuthenticationData) -> Self {
146            self.inner = self.inner.with_authentication_data(authentication);
147            self
148        }
149
150        /// Attaches a policy violation to this message.
151        pub fn with_policy_violation(mut self, violation: PolicyViolation) -> Self {
152            self.inner = self.inner.with_policy_violation(violation);
153            self
154        }
155    };
156}
157
158macro_rules! impl_unit_http_message {
159    ($type:ty) => {
160        impl UnitHttpMessage for $type {
161            fn header(&self, header: &str) -> Option<&str> {
162                self.inner.header(header)
163            }
164
165            fn headers(&self) -> &Vec<(String, String)> {
166                self.inner.headers()
167            }
168
169            fn body(&self) -> &[u8] {
170                self.inner.body()
171            }
172
173            fn property<K: Into<String>>(&self, key: Vec<K>) -> Option<Bytes> {
174                self.inner.property(key)
175            }
176
177            fn properties(&self) -> HashMap<Vec<String>, Bytes> {
178                self.inner.properties()
179            }
180
181            fn authentication(&self) -> Option<AuthenticationData> {
182                self.inner.authentication()
183            }
184
185            fn violation(&self) -> Option<PolicyViolation> {
186                self.inner.violation()
187            }
188        }
189    };
190}
191
192impl_unit_http_message!(UnitHttpRequest);
193impl_unit_http_message!(UnitHttpResponse);
194
195impl UnitHttpRequest {
196    /// Creates a request with the given HTTP method.
197    pub fn custom<M: Into<String>>(method: M) -> Self {
198        Self {
199            inner: RequestResponse::default().with_header(":method", method.into()),
200        }
201    }
202
203    http_method!(get, "GET");
204    http_method!(post, "POST");
205    http_method!(put, "PUT");
206    http_method!(patch, "PATCH");
207    http_method!(delete, "DELETE");
208    http_method!(head, "HEAD");
209    http_method!(options, "OPTIONS");
210
211    impl_request_response_methods!();
212
213    /// Sets the `:path` pseudo-header.
214    pub fn with_path<P: Into<String>>(mut self, path: P) -> Self {
215        self.inner = self.inner.with_header(":path", path.into());
216        self
217    }
218}
219
220impl From<RequestResponse> for UnitHttpRequest {
221    fn from(value: RequestResponse) -> Self {
222        Self { inner: value }
223    }
224}
225
226impl UnitHttpResponse {
227    /// Creates a response with the given HTTP status code.
228    pub fn new(status: u32) -> Self {
229        Self {
230            inner: RequestResponse::default().with_header(":status", status.to_string()),
231        }
232    }
233
234    impl_request_response_methods!();
235
236    /// Returns the HTTP status code parsed from the `:status` pseudo-header.
237    /// Returns `0` if the header is absent or not a valid integer.
238    pub fn status_code(&self) -> u32 {
239        self.inner
240            .header(":status")
241            .and_then(|s| s.parse().ok())
242            .unwrap_or_default()
243    }
244}
245
246impl From<RequestResponse> for UnitHttpResponse {
247    fn from(value: RequestResponse) -> Self {
248        Self { inner: value }
249    }
250}
251
252/// Represents an HTTP request or response in unit tests.
253///
254/// This struct is the primary data exchange object in `pdk-unit` tests, encapsulating
255/// HTTP headers, body, and a flexible properties map that can store Envoy-like metadata,
256/// authentication data, and policy violations.
257///
258/// # Example
259///
260/// ```ignore
261/// let request = RequestResponse::default()
262///     .with_header(":method", "GET")
263///     .with_header(":path", "/api/users")
264///     .with_header("authorization", "Bearer token123")
265///     .with_property(vec!["custom", "key"], "value");
266/// ```
267#[derive(Clone, Default, Serialize, Deserialize)]
268pub(crate) struct RequestResponse {
269    pub(crate) headers: Vec<(String, String)>,
270    pub(crate) body: Vec<u8>,
271    #[serde(deserialize_with = "de_properties")]
272    properties: Properties,
273}
274
275impl RequestResponse {
276    pub(crate) fn create(
277        headers: Vec<(String, String)>,
278        body: Vec<u8>,
279        properties: HashMap<Vec<String>, Bytes>,
280    ) -> Self {
281        Self {
282            headers,
283            body,
284            properties: Properties::new(properties),
285        }
286    }
287
288    /// Creates a new `RequestResponse` with the given headers and optional body.
289    ///
290    /// # Arguments
291    ///
292    /// * `headers` - A vector of key-value tuples representing HTTP headers
293    /// * `body` - An optional body payload
294    pub fn new(headers: Vec<(&str, &str)>, body: Option<&[u8]>) -> Self {
295        Self {
296            headers: headers
297                .into_iter()
298                .map(|(key, value)| (key.to_string(), value.into()))
299                .collect(),
300            body: body.map(|body| body.into()).unwrap_or_default(),
301            properties: Properties::default(),
302        }
303    }
304
305    /// Adds a header to the request/response. Returns `self` for method chaining.
306    pub fn with_header<K: Into<String>, V: Into<String>>(mut self, key: K, val: V) -> Self {
307        self.headers.push((key.into(), val.into()));
308        self
309    }
310
311    /// Returns the value of the specified header, or `None` if not found.
312    pub fn header(&self, header: &str) -> Option<&str> {
313        self.headers.iter().find_map(|(key, value)| {
314            if key.eq_ignore_ascii_case(header) {
315                Some(value.as_str())
316            } else {
317                None
318            }
319        })
320    }
321
322    /// Returns a reference to all headers.
323    pub fn headers(&self) -> &Vec<(String, String)> {
324        &self.headers
325    }
326
327    /// Sets the body of the request/response. Returns `self` for method chaining.
328    pub fn with_body<B: Into<Vec<u8>>>(mut self, body: B) -> Self {
329        self.body = body.into();
330        self
331    }
332
333    /// Returns a reference to the body as a byte slice.
334    pub fn body(&self) -> &[u8] {
335        self.body.as_slice()
336    }
337
338    /// Sets a property with a hierarchical key path. Returns `self` for method chaining.
339    ///
340    /// Properties are used to store Envoy-like metadata accessible via `get_property`.
341    pub fn with_property<K: Into<String>, V: Into<Vec<u8>>>(self, key: Vec<K>, value: V) -> Self {
342        let key = key.into_iter().map(|k| k.into()).collect();
343        self.properties
344            .properties
345            .borrow_mut()
346            .insert(key, value.into());
347        self
348    }
349
350    /// Returns the value of a property by its hierarchical key path, or `None` if not found.
351    pub fn property<K: Into<String>>(&self, key: Vec<K>) -> Option<Bytes> {
352        let key: Vec<String> = key.into_iter().map(|k| k.into()).collect();
353        self.properties.properties.borrow().get(&key).cloned()
354    }
355
356    /// Calling this method will override all properties, authentication data and policy violations
357    pub fn with_properties(mut self, properties: HashMap<Vec<String>, Bytes>) -> Self {
358        self.properties = Properties {
359            properties: Rc::new(RefCell::new(properties)),
360        };
361        self
362    }
363
364    /// Returns a clone of all properties as a `HashMap`.
365    pub fn properties(&self) -> HashMap<Vec<String>, Bytes> {
366        self.properties.properties.borrow().clone()
367    }
368
369    /// Sets authentication data on the request/response. Returns `self` for method chaining.
370    ///
371    /// This is used to simulate authenticated requests in tests.
372    pub fn with_authentication_data(self, authentication: AuthenticationData) -> Self {
373        Authentication::new(self.properties.shared()).set_authentication(Some(&authentication));
374        self
375    }
376
377    /// Returns the authentication data if present.
378    pub fn authentication(&self) -> Option<AuthenticationData> {
379        Authentication::new(self.properties.shared()).authentication()
380    }
381
382    /// Sets a policy violation on the request/response. Returns `self` for method chaining.
383    ///
384    /// This is used to simulate policy violation scenarios in tests.
385    pub fn with_policy_violation(self, violation: PolicyViolation) -> Self {
386        let violations = PolicyViolations::new(
387            self.properties.shared(),
388            violation.get_policy_name().to_string(),
389        );
390        if let Some(id) = violation.get_client_id() {
391            violations.generate_policy_violation_for_client_app(
392                violation.get_client_name().unwrap_or_default(),
393                id,
394            );
395        } else {
396            violations.generate_policy_violation();
397        }
398
399        self
400    }
401
402    /// Returns the policy violation if present.
403    pub fn violation(&self) -> Option<PolicyViolation> {
404        PolicyViolations::new(self.properties.shared(), String::default()).policy_violation()
405    }
406
407    pub(crate) fn with_property_if_missing<B: Into<String>>(
408        self,
409        key: &[&str],
410        bytes: B,
411    ) -> RequestResponse {
412        if self.property(key.to_vec()).is_none() {
413            self.with_property(key.to_vec(), bytes.into().into_bytes())
414        } else {
415            self
416        }
417    }
418}
419
420impl PartialEq for RequestResponse {
421    fn eq(&self, other: &Self) -> bool {
422        self.headers == other.headers && self.body == other.body
423    }
424}
425
426impl Debug for RequestResponse {
427    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
428        f.write_str("RequestResponse {")?;
429        f.write_str("headers: ")?;
430        f.write_str(format!("{:?}", self.headers).as_str())?;
431        f.write_str("body: ")?;
432        f.write_str(format!("{:?}", String::from_utf8_lossy(self.body.as_slice())).as_str())?;
433        f.write_str("}")
434    }
435}
436
437/// We create a custom struct to implement the Property accessor trait without exposing it to the user
438#[derive(Default)]
439struct Properties {
440    properties: Rc<RefCell<HashMap<Vec<String>, Bytes>>>,
441}
442
443impl Clone for Properties {
444    fn clone(&self) -> Self {
445        Self {
446            properties: Rc::new(RefCell::new(self.properties.borrow().clone())),
447        }
448    }
449}
450
451impl Properties {
452    pub fn new(properties: HashMap<Vec<String>, Bytes>) -> Self {
453        Self {
454            properties: Rc::new(RefCell::new(properties)),
455        }
456    }
457
458    pub fn shared(&self) -> Self {
459        Self {
460            properties: Rc::clone(&self.properties),
461        }
462    }
463}
464
465#[derive(Serialize, Deserialize)]
466struct Property {
467    key: Vec<String>,
468    value: Bytes,
469}
470
471impl Serialize for Properties {
472    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
473    where
474        S: Serializer,
475    {
476        self.properties
477            .borrow()
478            .clone()
479            .into_iter()
480            .map(|(key, value)| Property { key, value })
481            .collect::<Vec<Property>>()
482            .serialize(serializer)
483    }
484}
485
486fn de_properties<'de, D>(deserializer: D) -> Result<Properties, D::Error>
487where
488    D: Deserializer<'de>,
489{
490    let exp: Vec<Property> = serde::de::Deserialize::deserialize(deserializer)?;
491    Ok(Properties {
492        properties: Rc::new(RefCell::new(
493            exp.into_iter()
494                .map(|property| (property.key, property.value))
495                .collect(),
496        )),
497    })
498}
499
500impl PropertyAccessor for Properties {
501    fn read_property(&self, path: &[&str]) -> Option<Bytes> {
502        self.properties
503            .borrow()
504            .get(&path.iter().map(|s| s.to_string()).collect::<Vec<String>>())
505            .cloned()
506    }
507
508    fn set_property(&self, path: &[&str], value: Option<&[u8]>) {
509        match value {
510            None => {
511                self.properties
512                    .borrow_mut()
513                    .remove(&path.iter().map(|s| s.to_string()).collect::<Vec<String>>());
514            }
515            Some(value) => {
516                self.properties.borrow_mut().insert(
517                    path.iter().map(|s| s.to_string()).collect::<Vec<String>>(),
518                    value.into(),
519                );
520            }
521        }
522    }
523}
524
525/// Represents a gRPC request in unit tests.
526///
527/// This struct encapsulates the service name, method name, initial metadata,
528/// and the serialized protobuf message for gRPC calls made by policies.
529#[derive(Clone, Default, Serialize, Deserialize)]
530pub struct UnitGrpcRequest {
531    service: String,
532    method: String,
533    initial_metadata: Vec<(String, Bytes)>,
534    message: Option<Bytes>,
535}
536
537impl UnitGrpcRequest {
538    pub(crate) fn new(
539        service: &str,
540        method: &str,
541        initial_metadata: Vec<(&str, &[u8])>,
542        message: Option<&[u8]>,
543    ) -> UnitGrpcRequest {
544        UnitGrpcRequest {
545            service: service.to_string(),
546            method: method.to_string(),
547            initial_metadata: initial_metadata
548                .iter()
549                .map(|(key, value)| (key.to_string(), value.to_vec()))
550                .collect(),
551            message: message.map(|m| m.to_vec()),
552        }
553    }
554
555    /// Returns the gRPC service name.
556    pub fn service(&self) -> &str {
557        &self.service
558    }
559
560    /// Returns the gRPC method name.
561    pub fn method(&self) -> &str {
562        &self.method
563    }
564
565    /// Returns a reference to the initial metadata (headers) of the gRPC request.
566    pub fn initial_metadata(&self) -> &Vec<(String, Bytes)> {
567        &self.initial_metadata
568    }
569
570    /// Returns a reference to the serialized protobuf message, if present.
571    pub fn message(&self) -> Option<&Bytes> {
572        self.message.as_ref()
573    }
574}
575
576/// Represents a gRPC response in unit tests.
577///
578/// This struct is used to mock gRPC responses from backend services.
579///
580/// # Example
581///
582/// ```ignore
583/// use pdk_unit::UnitGrpcResponse;
584///
585/// let response = UnitGrpcResponse::default()
586///     .with_status_code(0)
587///     .with_message(serialized_protobuf);
588/// ```
589#[derive(Clone, Default, Serialize, Deserialize)]
590pub struct UnitGrpcResponse {
591    pub(crate) status_code: u32,
592    pub(crate) status: Option<String>,
593    pub(crate) message: Bytes,
594}
595
596impl UnitGrpcResponse {
597    /// Sets the gRPC status code. Returns `self` for method chaining.
598    ///
599    /// A status code of `0` indicates success (OK).
600    pub fn with_status_code(mut self, status: u32) -> Self {
601        self.status_code = status;
602        self
603    }
604
605    /// Sets the serialized protobuf response message. Returns `self` for method chaining.
606    pub fn with_message(mut self, message: Vec<u8>) -> Self {
607        self.message = message;
608        self
609    }
610
611    /// Sets the gRPC status message. Returns `self` for method chaining.
612    pub fn with_status<S: Into<String>>(mut self, status: S) -> Self {
613        self.status = Some(status.into());
614        self
615    }
616}
pdk_unit/tester/io.rs

pdk_unit/tester/
io.rs
