Skip to main content

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