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

pdk_unit/tester/
io.rs
