pdk_unit/tester/
builder.rs

1// Copyright (c) 2026, Salesforce, Inc.,
2// All rights reserved.
3// For full license text, see the LICENSE.txt file
4#[cfg(feature = "pdk")]
5use crate::tester::unit_test::IDENTITY_MANAGEMENT_SVC;
6use crate::tester::unit_test::{Backends, UnitTest, UnitTestConfig};
7use crate::{Backend, GrpcBackend, TraceBackend, UnitHttpResponse};
8#[cfg(feature = "pdk")]
9use classy::Entrypoint;
10#[cfg(feature = "pdk")]
11use pdk_core::client::service_name;
12#[cfg(feature = "proxy-wasm-rust-sdk")]
13use proxy_wasm_stub::traits::RootContext;
14use std::rc::Rc;
15
16/// A fluent builder for configuring and creating [`UnitTest`] instances.
17///
18/// This is the entry point for setting up unit tests. It allows configuring the policy
19/// under test, its configuration, and mocking external dependencies like HTTP and gRPC backends.
20///
21/// # Example
22///
23/// ```ignore
24/// let mut tester = UnitTestBuilder::default()
25///     .with_config(json!({"key": "value"}).to_string())
26///     .with_backend(TraceBackend::new(UnitHttpResponse::new(200)))
27///     .with_http_upstream_from_authority("api.example.com", my_backend)
28///     .with_entrypoint(crate::configure);
29/// ```
30pub struct UnitTestBuilder {
31    backend: Backends,
32    config: UnitTestConfig,
33}
34
35impl Default for UnitTestBuilder {
36    fn default() -> Self {
37        Self {
38            backend: Backends {
39                backend: Box::new(TraceBackend::new(UnitHttpResponse::new(200))),
40                upstreams: Default::default(),
41                grpc_upstreams: Default::default(),
42            },
43            config: Default::default(),
44        }
45    }
46}
47
48impl UnitTestBuilder {
49    /// Sets the policy configuration JSON. Returns `self` for method chaining.
50    pub fn with_config<S: Into<String>>(mut self, config: S) -> Self {
51        self.config.policy_config = config.into();
52        self
53    }
54
55    /// Sets the default HTTP backend used when no specific upstream is matched.
56    /// Returns `self` for method chaining.
57    pub fn with_backend<B: Backend + 'static>(mut self, backend: B) -> Self {
58        self.backend.backend = Box::new(backend);
59        self
60    }
61
62    /// Registers an HTTP upstream backend by its full service name.
63    /// Returns `self` for method chaining.
64    pub fn with_http_upstream<K: Into<String>, B: Backend + 'static>(
65        mut self,
66        upstream: K,
67        backend: B,
68    ) -> Self {
69        self.backend
70            .upstreams
71            .insert(upstream.into(), Rc::new(backend));
72        self
73    }
74
75    /// Registers a gRPC upstream backend by its full service name.
76    /// Returns `self` for method chaining.
77    pub fn with_grpc_upstream<K: Into<String>, B: GrpcBackend + 'static>(
78        mut self,
79        upstream: K,
80        backend: B,
81    ) -> Self {
82        self.backend
83            .grpc_upstreams
84            .insert(upstream.into(), Rc::new(backend));
85        self
86    }
87}
88
89#[cfg(feature = "pdk")]
90impl UnitTestBuilder {
91    /// Registers an HTTP upstream backend by authority (e.g., "api.example.com").
92    ///
93    /// The authority is automatically converted to the internal upstream service name format.
94    /// Returns `self` for method chaining.
95    pub fn with_http_upstream_from_authority<K: AsRef<str>, B: Backend + 'static>(
96        self,
97        authority: K,
98        backend: B,
99    ) -> Self {
100        let upstream = self.upstream(authority.as_ref());
101        self.with_http_upstream(upstream, backend)
102    }
103
104    /// Registers a gRPC upstream backend by authority (e.g., "grpc.example.com").
105    ///
106    /// The authority is automatically converted to the internal upstream service name format.
107    /// Returns `self` for method chaining.
108    pub fn with_grpc_upstream_from_authority<K: AsRef<str>, B: GrpcBackend + 'static>(
109        self,
110        authority: K,
111        backend: B,
112    ) -> Self {
113        let upstream = self.upstream(authority.as_ref());
114        self.with_grpc_upstream(upstream, backend)
115    }
116
117    /// Returns a mutable reference to the metadata that will be used for the test.
118    /// Modify the values as desired.
119    /// Modifying the values impact in the service registration process, if you'll use this method
120    /// make sure to do it before calling the `with_*_upstream_from_authority` methods.
121    pub fn metadata<F: FnOnce(&mut pdk_core::policy_context::api::Metadata)>(
122        mut self,
123        function: F,
124    ) -> Self {
125        function(&mut self.config.metadata);
126        self
127    }
128
129    #[cfg(feature = "experimental_local_mode")]
130    pub fn local_mode(mut self) -> Self {
131        self.config.local_mode = true;
132        self
133    }
134
135    /// Sets the function to be used as the identity management backend.
136    /// Returns `self` for method chaining.
137    pub fn with_identity_management<K: Into<String>, B: Backend + 'static>(
138        mut self,
139        url: K,
140        backend: B,
141    ) -> Self {
142        self.config.identity_management = Some(url.into());
143        self.with_http_upstream(IDENTITY_MANAGEMENT_SVC, backend)
144    }
145
146    /// Finalizes the builder and creates a [`UnitTest`] instance with the specified policy entrypoint.
147    ///
148    /// This method consumes the builder and initializes the test environment.
149    pub fn with_entrypoint<C, T, E: Entrypoint<C, T> + Clone + 'static>(
150        self,
151        entrypoint: E,
152    ) -> UnitTest {
153        UnitTest::new(entrypoint, self.config, self.backend)
154    }
155
156    fn upstream(&self, authority: &str) -> String {
157        format!(
158            "{}.{}.svc",
159            service_name(
160                self.config.metadata.policy_metadata.policy_name.as_str(),
161                authority
162            ),
163            self.config
164                .metadata
165                .policy_metadata
166                .policy_namespace
167                .as_str(),
168        )
169    }
170}
171
172#[cfg(feature = "proxy-wasm-rust-sdk")]
173impl UnitTestBuilder {
174    /// Finalizes the builder and creates a [`UnitTest`] instance with a raw proxy-wasm
175    /// [`RootContext`] factory. Use this when not using the PDK framework.
176    pub fn with_context<F>(self, factory: F) -> UnitTest
177    where
178        F: Fn() -> Box<dyn RootContext> + 'static,
179    {
180        UnitTest::new_with_context(factory, self.config, self.backend)
181    }
182}
pdk_unit/tester/builder.rs

pdk_unit/tester/
builder.rs
