pdk-unit 1.8.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use classy::stream::PropertyAccessor;
use pdk_core::policy_context::authentication::{
    Authentication, AuthenticationData, AuthenticationHandler,
};
use pdk_core::policy_context::policy_violation::{PolicyViolation, PolicyViolations};
use proxy_wasm_stub::types::Bytes;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::cell::RefCell;
use std::collections::HashMap;
use std::fmt::{Debug, Formatter};
use std::rc::Rc;

/// An HTTP request used in `pdk-unit` tests.
///
/// Construct one using the HTTP method constructors ([`get`](Self::get), [`post`](Self::post), etc.)
/// or [`custom`](Self::custom) for non-standard methods, then chain `with_*` builder methods to
/// populate headers, body, properties, and authentication.
///
/// Read-only accessors are available via the [`UnitHttpMessage`] trait.
///
/// # Example
///
/// ```ignore
/// use pdk_unit::{UnitHttpRequest, UnitHttpMessage};
///
/// let req = UnitHttpRequest::get()
///     .with_path("/api/users")
///     .with_header("authorization", "Bearer token123")
///     .with_body("hello");
///
/// assert_eq!(req.header("authorization"), Some("Bearer token123"));
/// ```
#[derive(Clone, Debug, PartialEq)]
pub struct UnitHttpRequest {
    pub(crate) inner: RequestResponse,
}

/// An HTTP response used in `pdk-unit` tests.
///
/// Construct one with [`new`](Self::new) providing the HTTP status code, then chain `with_*`
/// builder methods to populate headers and body.
///
/// Read-only accessors are available via the [`UnitHttpMessage`] trait.
///
/// # Example
///
/// ```ignore
/// use pdk_unit::{UnitHttpResponse, UnitHttpMessage};
///
/// let resp = UnitHttpResponse::new(200)
///     .with_header("content-type", "application/json")
///     .with_body(r#"{"ok":true}"#);
///
/// assert_eq!(resp.status_code(), 200);
/// assert_eq!(resp.header("content-type"), Some("application/json"));
/// ```
#[derive(Clone, Debug, PartialEq)]
pub struct UnitHttpResponse {
    pub(crate) inner: RequestResponse,
}

macro_rules! http_method {
    ($fn_name:ident, $method:expr) => {
        #[doc = "Creates a `"]
        #[doc = $method]
        #[doc = "` request."]
        pub fn $fn_name() -> Self {
            Self::custom($method)
        }
    };
}

/// Common read-only accessors for HTTP messages in unit tests.
///
/// Implemented by both [`UnitHttpRequest`] and [`UnitHttpResponse`], allowing test helper
/// functions to be written generically over either type.
///
/// # Example
///
/// ```ignore
/// use pdk_unit::{UnitHttpMessage, UnitHttpRequest, UnitHttpResponse};
///
/// fn assert_ok<M: UnitHttpMessage>(msg: &M) {
///     assert!(msg.body().len() > 0);
/// }
/// ```
pub trait UnitHttpMessage {
    /// Returns the value of the header with the given name, or `None` if not present.
    fn header(&self, header: &str) -> Option<&str>;

    /// Returns all headers as a list of `(name, value)` pairs.
    fn headers(&self) -> &Vec<(String, String)>;

    /// Returns the message body as a byte slice.
    fn body(&self) -> &[u8];

    /// Returns the value of a property identified by the given key path, or `None` if not set.
    fn property<K: Into<String>>(&self, key: Vec<K>) -> Option<Bytes>;

    /// Returns all properties as a map of key path to raw bytes.
    fn properties(&self) -> HashMap<Vec<String>, Bytes>;

    /// Returns the authentication data attached to this message, if any.
    fn authentication(&self) -> Option<AuthenticationData>;

    /// Returns the policy violation attached to this message, if any.
    fn violation(&self) -> Option<PolicyViolation>;
}

macro_rules! impl_request_response_methods {
    () => {
        /// Sets a header, replacing any existing value for the same name.
        pub fn with_header<K: Into<String>, V: Into<String>>(mut self, key: K, val: V) -> Self {
            self.inner = self.inner.with_header(key, val);
            self
        }

        /// Sets the message body.
        pub fn with_body<B: Into<Vec<u8>>>(mut self, body: B) -> Self {
            self.inner = self.inner.with_body(body);
            self
        }

        /// Sets a single property identified by the given key path.
        pub fn with_property<K: Into<String>, V: Into<Vec<u8>>>(
            mut self,
            key: Vec<K>,
            value: V,
        ) -> Self {
            self.inner = self.inner.with_property(key, value);
            self
        }

        /// Replaces all properties with the given map.
        pub fn with_properties(mut self, properties: HashMap<Vec<String>, Bytes>) -> Self {
            self.inner = self.inner.with_properties(properties);
            self
        }

        /// Attaches authentication data to this message.
        pub fn with_authentication_data(mut self, authentication: AuthenticationData) -> Self {
            self.inner = self.inner.with_authentication_data(authentication);
            self
        }

        /// Attaches a policy violation to this message.
        pub fn with_policy_violation(mut self, violation: PolicyViolation) -> Self {
            self.inner = self.inner.with_policy_violation(violation);
            self
        }
    };
}

macro_rules! impl_unit_http_message {
    ($type:ty) => {
        impl UnitHttpMessage for $type {
            fn header(&self, header: &str) -> Option<&str> {
                self.inner.header(header)
            }

            fn headers(&self) -> &Vec<(String, String)> {
                self.inner.headers()
            }

            fn body(&self) -> &[u8] {
                self.inner.body()
            }

            fn property<K: Into<String>>(&self, key: Vec<K>) -> Option<Bytes> {
                self.inner.property(key)
            }

            fn properties(&self) -> HashMap<Vec<String>, Bytes> {
                self.inner.properties()
            }

            fn authentication(&self) -> Option<AuthenticationData> {
                self.inner.authentication()
            }

            fn violation(&self) -> Option<PolicyViolation> {
                self.inner.violation()
            }
        }
    };
}

impl_unit_http_message!(UnitHttpRequest);
impl_unit_http_message!(UnitHttpResponse);

impl UnitHttpRequest {
    /// Creates a request with the given HTTP method.
    pub fn custom<M: Into<String>>(method: M) -> Self {
        Self {
            inner: RequestResponse::default().with_header(":method", method.into()),
        }
    }

    http_method!(get, "GET");
    http_method!(post, "POST");
    http_method!(put, "PUT");
    http_method!(patch, "PATCH");
    http_method!(delete, "DELETE");
    http_method!(head, "HEAD");
    http_method!(options, "OPTIONS");

    impl_request_response_methods!();

    /// Sets the `:path` pseudo-header.
    pub fn with_path<P: Into<String>>(mut self, path: P) -> Self {
        self.inner = self.inner.with_header(":path", path.into());
        self
    }
}

impl From<RequestResponse> for UnitHttpRequest {
    fn from(value: RequestResponse) -> Self {
        Self { inner: value }
    }
}

impl UnitHttpResponse {
    /// Creates a response with the given HTTP status code.
    pub fn new(status: u32) -> Self {
        Self {
            inner: RequestResponse::default().with_header(":status", status.to_string()),
        }
    }

    impl_request_response_methods!();

    /// Returns the HTTP status code parsed from the `:status` pseudo-header.
    /// Returns `0` if the header is absent or not a valid integer.
    pub fn status_code(&self) -> u32 {
        self.inner
            .header(":status")
            .and_then(|s| s.parse().ok())
            .unwrap_or_default()
    }
}

impl From<RequestResponse> for UnitHttpResponse {
    fn from(value: RequestResponse) -> Self {
        Self { inner: value }
    }
}

/// Represents an HTTP request or response in unit tests.
///
/// This struct is the primary data exchange object in `pdk-unit` tests, encapsulating
/// HTTP headers, body, and a flexible properties map that can store Envoy-like metadata,
/// authentication data, and policy violations.
///
/// # Example
///
/// ```ignore
/// let request = RequestResponse::default()
///     .with_header(":method", "GET")
///     .with_header(":path", "/api/users")
///     .with_header("authorization", "Bearer token123")
///     .with_property(vec!["custom", "key"], "value");
/// ```
#[derive(Clone, Default, Serialize, Deserialize)]
pub(crate) struct RequestResponse {
    pub(crate) headers: Vec<(String, String)>,
    pub(crate) body: Vec<u8>,
    #[serde(deserialize_with = "de_properties")]
    properties: Properties,
}

impl RequestResponse {
    pub(crate) fn create(
        headers: Vec<(String, String)>,
        body: Vec<u8>,
        properties: HashMap<Vec<String>, Bytes>,
    ) -> Self {
        Self {
            headers,
            body,
            properties: Properties::new(properties),
        }
    }

    /// Creates a new `RequestResponse` with the given headers and optional body.
    ///
    /// # Arguments
    ///
    /// * `headers` - A vector of key-value tuples representing HTTP headers
    /// * `body` - An optional body payload
    pub fn new(headers: Vec<(&str, &str)>, body: Option<&[u8]>) -> Self {
        Self {
            headers: headers
                .into_iter()
                .map(|(key, value)| (key.to_string(), value.into()))
                .collect(),
            body: body.map(|body| body.into()).unwrap_or_default(),
            properties: Properties::default(),
        }
    }

    /// Adds a header to the request/response. Returns `self` for method chaining.
    pub fn with_header<K: Into<String>, V: Into<String>>(mut self, key: K, val: V) -> Self {
        self.headers.push((key.into(), val.into()));
        self
    }

    /// Returns the value of the specified header, or `None` if not found.
    pub fn header(&self, header: &str) -> Option<&str> {
        self.headers.iter().find_map(|(key, value)| {
            if key.eq_ignore_ascii_case(header) {
                Some(value.as_str())
            } else {
                None
            }
        })
    }

    /// Returns a reference to all headers.
    pub fn headers(&self) -> &Vec<(String, String)> {
        &self.headers
    }

    /// Sets the body of the request/response. Returns `self` for method chaining.
    pub fn with_body<B: Into<Vec<u8>>>(mut self, body: B) -> Self {
        self.body = body.into();
        self
    }

    /// Returns a reference to the body as a byte slice.
    pub fn body(&self) -> &[u8] {
        self.body.as_slice()
    }

    /// Sets a property with a hierarchical key path. Returns `self` for method chaining.
    ///
    /// Properties are used to store Envoy-like metadata accessible via `get_property`.
    pub fn with_property<K: Into<String>, V: Into<Vec<u8>>>(self, key: Vec<K>, value: V) -> Self {
        let key = key.into_iter().map(|k| k.into()).collect();
        self.properties
            .properties
            .borrow_mut()
            .insert(key, value.into());
        self
    }

    /// Returns the value of a property by its hierarchical key path, or `None` if not found.
    pub fn property<K: Into<String>>(&self, key: Vec<K>) -> Option<Bytes> {
        let key: Vec<String> = key.into_iter().map(|k| k.into()).collect();
        self.properties.properties.borrow().get(&key).cloned()
    }

    /// Calling this method will override all properties, authentication data and policy violations
    pub fn with_properties(mut self, properties: HashMap<Vec<String>, Bytes>) -> Self {
        self.properties = Properties {
            properties: Rc::new(RefCell::new(properties)),
        };
        self
    }

    /// Returns a clone of all properties as a `HashMap`.
    pub fn properties(&self) -> HashMap<Vec<String>, Bytes> {
        self.properties.properties.borrow().clone()
    }

    /// Sets authentication data on the request/response. Returns `self` for method chaining.
    ///
    /// This is used to simulate authenticated requests in tests.
    pub fn with_authentication_data(self, authentication: AuthenticationData) -> Self {
        Authentication::new(self.properties.shared()).set_authentication(Some(&authentication));
        self
    }

    /// Returns the authentication data if present.
    pub fn authentication(&self) -> Option<AuthenticationData> {
        Authentication::new(self.properties.shared()).authentication()
    }

    /// Sets a policy violation on the request/response. Returns `self` for method chaining.
    ///
    /// This is used to simulate policy violation scenarios in tests.
    pub fn with_policy_violation(self, violation: PolicyViolation) -> Self {
        let violations = PolicyViolations::new(
            self.properties.shared(),
            violation.get_policy_name().to_string(),
        );
        if let Some(id) = violation.get_client_id() {
            violations.generate_policy_violation_for_client_app(
                violation.get_client_name().unwrap_or_default(),
                id,
            );
        } else {
            violations.generate_policy_violation();
        }

        self
    }

    /// Returns the policy violation if present.
    pub fn violation(&self) -> Option<PolicyViolation> {
        PolicyViolations::new(self.properties.shared(), String::default()).policy_violation()
    }

    pub(crate) fn with_property_if_missing<B: Into<String>>(
        self,
        key: &[&str],
        bytes: B,
    ) -> RequestResponse {
        if self.property(key.to_vec()).is_none() {
            self.with_property(key.to_vec(), bytes.into().into_bytes())
        } else {
            self
        }
    }
}

impl PartialEq for RequestResponse {
    fn eq(&self, other: &Self) -> bool {
        self.headers == other.headers && self.body == other.body
    }
}

impl Debug for RequestResponse {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str("RequestResponse {")?;
        f.write_str("headers: ")?;
        f.write_str(format!("{:?}", self.headers).as_str())?;
        f.write_str("body: ")?;
        f.write_str(format!("{:?}", String::from_utf8_lossy(self.body.as_slice())).as_str())?;
        f.write_str("}")
    }
}

/// We create a custom struct to implement the Property accessor trait without exposing it to the user
#[derive(Default)]
struct Properties {
    properties: Rc<RefCell<HashMap<Vec<String>, Bytes>>>,
}

impl Clone for Properties {
    fn clone(&self) -> Self {
        Self {
            properties: Rc::new(RefCell::new(self.properties.borrow().clone())),
        }
    }
}

impl Properties {
    pub fn new(properties: HashMap<Vec<String>, Bytes>) -> Self {
        Self {
            properties: Rc::new(RefCell::new(properties)),
        }
    }

    pub fn shared(&self) -> Self {
        Self {
            properties: Rc::clone(&self.properties),
        }
    }
}

#[derive(Serialize, Deserialize)]
struct Property {
    key: Vec<String>,
    value: Bytes,
}

impl Serialize for Properties {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        self.properties
            .borrow()
            .clone()
            .into_iter()
            .map(|(key, value)| Property { key, value })
            .collect::<Vec<Property>>()
            .serialize(serializer)
    }
}

fn de_properties<'de, D>(deserializer: D) -> Result<Properties, D::Error>
where
    D: Deserializer<'de>,
{
    let exp: Vec<Property> = serde::de::Deserialize::deserialize(deserializer)?;
    Ok(Properties {
        properties: Rc::new(RefCell::new(
            exp.into_iter()
                .map(|property| (property.key, property.value))
                .collect(),
        )),
    })
}

impl PropertyAccessor for Properties {
    fn read_property(&self, path: &[&str]) -> Option<Bytes> {
        self.properties
            .borrow()
            .get(&path.iter().map(|s| s.to_string()).collect::<Vec<String>>())
            .cloned()
    }

    fn set_property(&self, path: &[&str], value: Option<&[u8]>) {
        match value {
            None => {
                self.properties
                    .borrow_mut()
                    .remove(&path.iter().map(|s| s.to_string()).collect::<Vec<String>>());
            }
            Some(value) => {
                self.properties.borrow_mut().insert(
                    path.iter().map(|s| s.to_string()).collect::<Vec<String>>(),
                    value.into(),
                );
            }
        }
    }
}

/// Represents a gRPC request in unit tests.
///
/// This struct encapsulates the service name, method name, initial metadata,
/// and the serialized protobuf message for gRPC calls made by policies.
#[derive(Clone, Default, Serialize, Deserialize)]
pub struct UnitGrpcRequest {
    service: String,
    method: String,
    initial_metadata: Vec<(String, Bytes)>,
    message: Option<Bytes>,
}

impl UnitGrpcRequest {
    pub(crate) fn new(
        service: &str,
        method: &str,
        initial_metadata: Vec<(&str, &[u8])>,
        message: Option<&[u8]>,
    ) -> UnitGrpcRequest {
        UnitGrpcRequest {
            service: service.to_string(),
            method: method.to_string(),
            initial_metadata: initial_metadata
                .iter()
                .map(|(key, value)| (key.to_string(), value.to_vec()))
                .collect(),
            message: message.map(|m| m.to_vec()),
        }
    }

    /// Returns the gRPC service name.
    pub fn service(&self) -> &str {
        &self.service
    }

    /// Returns the gRPC method name.
    pub fn method(&self) -> &str {
        &self.method
    }

    /// Returns a reference to the initial metadata (headers) of the gRPC request.
    pub fn initial_metadata(&self) -> &Vec<(String, Bytes)> {
        &self.initial_metadata
    }

    /// Returns a reference to the serialized protobuf message, if present.
    pub fn message(&self) -> Option<&Bytes> {
        self.message.as_ref()
    }
}

/// Represents a gRPC response in unit tests.
///
/// This struct is used to mock gRPC responses from backend services.
///
/// # Example
///
/// ```ignore
/// use pdk_unit::UnitGrpcResponse;
///
/// let response = UnitGrpcResponse::default()
///     .with_status_code(0)
///     .with_message(serialized_protobuf);
/// ```
#[derive(Clone, Default, Serialize, Deserialize)]
pub struct UnitGrpcResponse {
    pub(crate) status_code: u32,
    pub(crate) status: Option<String>,
    pub(crate) message: Bytes,
}

impl UnitGrpcResponse {
    /// Sets the gRPC status code. Returns `self` for method chaining.
    ///
    /// A status code of `0` indicates success (OK).
    pub fn with_status_code(mut self, status: u32) -> Self {
        self.status_code = status;
        self
    }

    /// Sets the serialized protobuf response message. Returns `self` for method chaining.
    pub fn with_message(mut self, message: Vec<u8>) -> Self {
        self.message = message;
        self
    }

    /// Sets the gRPC status message. Returns `self` for method chaining.
    pub fn with_status<S: Into<String>>(mut self, status: S) -> Self {
        self.status = Some(status.into());
        self
    }
}