pdk-unit 1.8.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file
use crate::tester::unit_test::{Backends, UnitTest, UnitTestConfig, IDENTITY_MANAGEMENT_SVC};
use crate::{Backend, GrpcBackend, TraceBackend, UnitHttpResponse};
use classy::Entrypoint;
use pdk_core::client::service_name;
use std::rc::Rc;

/// A fluent builder for configuring and creating [`UnitTest`] instances.
///
/// This is the entry point for setting up unit tests. It allows configuring the policy
/// under test, its configuration, and mocking external dependencies like HTTP and gRPC backends.
///
/// # Example
///
/// ```ignore
/// let mut tester = UnitTestBuilder::default()
///     .with_config(json!({"key": "value"}).to_string())
///     .with_backend(TraceBackend::new(UnitHttpResponse::new(200)))
///     .with_http_upstream_from_authority("api.example.com", my_backend)
///     .with_entrypoint(crate::configure);
/// ```
pub struct UnitTestBuilder {
    backend: Backends,
    config: UnitTestConfig,
}

impl Default for UnitTestBuilder {
    fn default() -> Self {
        Self {
            backend: Backends {
                backend: Box::new(TraceBackend::new(UnitHttpResponse::new(200))),
                upstreams: Default::default(),
                grpc_upstreams: Default::default(),
            },
            config: Default::default(),
        }
    }
}

impl UnitTestBuilder {
    /// Sets the policy configuration JSON. Returns `self` for method chaining.
    pub fn with_config<S: Into<String>>(mut self, config: S) -> Self {
        self.config.policy_config = config.into();
        self
    }

    /// Sets the default HTTP backend used when no specific upstream is matched.
    /// Returns `self` for method chaining.
    pub fn with_backend<B: Backend + 'static>(mut self, backend: B) -> Self {
        self.backend.backend = Box::new(backend);
        self
    }

    /// Registers an HTTP upstream backend by authority (e.g., "api.example.com").
    ///
    /// The authority is automatically converted to the internal upstream service name format.
    /// Returns `self` for method chaining.
    pub fn with_http_upstream_from_authority<K: AsRef<str>, B: Backend + 'static>(
        self,
        authority: K,
        backend: B,
    ) -> Self {
        let upstream = self.upstream(authority.as_ref());
        self.with_http_upstream(upstream, backend)
    }

    /// Registers an HTTP upstream backend by its full service name.
    /// Returns `self` for method chaining.
    pub fn with_http_upstream<K: Into<String>, B: Backend + 'static>(
        mut self,
        upstream: K,
        backend: B,
    ) -> Self {
        self.backend
            .upstreams
            .insert(upstream.into(), Rc::new(backend));
        self
    }

    /// Registers a gRPC upstream backend by authority (e.g., "grpc.example.com").
    ///
    /// The authority is automatically converted to the internal upstream service name format.
    /// Returns `self` for method chaining.
    pub fn with_grpc_upstream_from_authority<K: AsRef<str>, B: GrpcBackend + 'static>(
        self,
        authority: K,
        backend: B,
    ) -> Self {
        let upstream = self.upstream(authority.as_ref());
        self.with_grpc_upstream(upstream, backend)
    }

    /// Registers a gRPC upstream backend by its full service name.
    /// Returns `self` for method chaining.
    pub fn with_grpc_upstream<K: Into<String>, B: GrpcBackend + 'static>(
        mut self,
        upstream: K,
        backend: B,
    ) -> Self {
        self.backend
            .grpc_upstreams
            .insert(upstream.into(), Rc::new(backend));
        self
    }

    /// Returns a mutable reference to the metadata that will be used for the test.
    /// Modify the values as desired.
    /// Modifying the values impact in the service registration process, if you'll use this method
    /// make sure to do it before calling the `with_*_upstream_from_authority` methods.
    pub fn metadata<F: FnOnce(&mut pdk_core::policy_context::api::Metadata)>(
        mut self,
        function: F,
    ) -> Self {
        function(&mut self.config.metadata);
        self
    }

    /// Sets the function to be used as the identity management backend.
    /// Returns `self` for method chaining.
    pub fn with_identity_management<K: Into<String>, B: Backend + 'static>(
        mut self,
        url: K,
        backend: B,
    ) -> Self {
        self.config.identity_management = Some(url.into());
        self.with_http_upstream(IDENTITY_MANAGEMENT_SVC, backend)
    }

    /// Finalizes the builder and creates a [`UnitTest`] instance with the specified policy entrypoint.
    ///
    /// This method consumes the builder and initializes the test environment.
    pub fn with_entrypoint<C, T, E: Entrypoint<C, T> + Clone + 'static>(
        self,
        entrypoint: E,
    ) -> UnitTest {
        UnitTest::new(entrypoint, self.config, self.backend)
    }

    fn upstream(&self, authority: &str) -> String {
        format!(
            "{}.{}.svc",
            service_name(
                self.config.metadata.policy_metadata.policy_name.as_str(),
                authority
            ),
            self.config
                .metadata
                .policy_metadata
                .policy_namespace
                .as_str(),
        )
    }
}