pdk_test/

composite.rs

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

//! Test composite orchestration
//!
//! This module provides the main test environment orchestration functionality.
//! It manages Docker containers, networks, and services in a coordinated way
//! for integration testing scenarios.
//!
//! ## Primary types
//!
//! - [`TestComposite`]: orchestrates test environment lifecycle
//! - [`TestCompositeBuilder`]: configures and builds test environments
//!

use std::any::{type_name, Any, TypeId};
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::rc::Rc;
use std::time::Duration;

use crate::cleanup::Cleanup;
use bollard::Docker;
use futures::future::{join_all, try_join_all};
use tokio::runtime::{Handle, RuntimeFlavor};

use crate::config::Config;
use crate::container::Container;
use crate::error::TestError;
use crate::host::Host;
use crate::network::Network;
use crate::runner::Test;
use crate::service::Service;
use crate::services::httpmock::HttpMockConfig;

struct UntypedConfig {
    erased: Rc<dyn Any>,
    source: Rc<dyn Config>,
}

impl UntypedConfig {
    fn new<T: Config + 'static>(config: T) -> Self {
        let erased: Rc<dyn Any> = Rc::new(config);
        let source: Rc<dyn Config> = erased.clone().downcast::<T>().unwrap();
        Self { erased, source }
    }

    fn upcast(&self) -> &dyn Config {
        self.source.as_ref()
    }

    fn downcast<T: Config + 'static>(&self) -> &T {
        self.erased.downcast_ref().unwrap()
    }
}

struct Inner {
    configs: HashMap<TypeId, HashMap<String, UntypedConfig>>,
    containers: HashMap<String, Container>,
    network: Network,
    test: Rc<Test>,
}

impl Inner {
    fn configs<T: Service + 'static>(&self) -> Result<&HashMap<String, UntypedConfig>, TestError> {
        self.configs
            .get(&TypeId::of::<T::Config>())
            .ok_or(TestError::UnknownService(type_name::<T>()))
    }

    fn service<T: Service + 'static>(&self) -> Result<T, TestError> {
        let config = self.configs::<T>()?.values().next().unwrap();
        let container = self.containers.get(config.upcast().hostname()).unwrap();
        Ok(T::new(config.downcast(), container))
    }

    fn service_by_hostname<T: Service + 'static>(&self, hostname: &str) -> Result<T, TestError> {
        let config = self
            .configs::<T>()?
            .get(hostname)
            .ok_or_else(|| TestError::UnknownServiceHostname(hostname.to_string()))?;
        let container = self.containers.get(config.upcast().hostname()).unwrap();
        Ok(T::new(config.downcast(), container))
    }
}

/// Main test environment orchestrator for integration tests.
pub struct TestComposite {
    inner: Option<Inner>,
}

impl TestComposite {
    /// Creates a new builder for configuring test environments.
    pub fn builder() -> TestCompositeBuilder {
        TestCompositeBuilder::new()
    }

    /// Gets a service instance by its type.
    /// Returns the first configured service of the specified type.
    pub fn service<T: Service + 'static>(&self) -> Result<T, TestError> {
        self.inner().service()
    }

    /// Gets a service instance by its hostname.
    /// Returns the service with the specified hostname and type.
    pub fn service_by_hostname<T: Service + 'static>(&self, name: &str) -> Result<T, TestError> {
        self.inner().service_by_hostname(name)
    }

    fn inner(&self) -> &Inner {
        self.inner.as_ref().unwrap()
    }
}

fn check_runtime() -> Result<(), TestError> {
    let handle = Handle::try_current()?;
    if !matches!(handle.runtime_flavor(), RuntimeFlavor::MultiThread) {
        return Err(TestError::UnavailableMultiThread);
    }
    Ok(())
}

async fn check_docker(docker: &Docker) -> Result<(), TestError> {
    match docker.ping().await {
        Ok(_) => Ok(()),
        Err(err) => Err(TestError::UnavailableDocker(err.into())),
    }
}

/// Builder for configuring and creating test environments.
///
/// This builder is used to configure the test environment with the services
/// that will be used in the test. It provides an API for adding services
/// and configuring them.
pub struct TestCompositeBuilder {
    configs: HashMap<TypeId, HashMap<String, UntypedConfig>>,
}

impl TestCompositeBuilder {
    fn new() -> Self {
        Self {
            configs: HashMap::new(),
        }
    }

    /// Configures a service with the provided configuration. The service will be
    /// started when the test environment is built. Each service must have a unique hostname.
    pub fn with_service<C: Config + 'static>(mut self, config: C) -> Self {
        let entry = self
            .configs
            .entry(TypeId::of::<C>())
            .or_default()
            .entry(config.hostname().to_string());

        match entry {
            Entry::Occupied(_) => panic!("Name {} configured twice", config.hostname()),
            Entry::Vacant(e) => e.insert(UntypedConfig::new(config)),
        };

        self
    }

    /// Builds the test environment with all configured services.
    /// Starts Docker containers, creates networks and initializes all services.
    /// Returns a `TestComposite` that can be used to access the services.
    pub async fn build(self) -> Result<TestComposite, TestError> {
        check_runtime()?;

        let httpmock_configs_len = self
            .configs
            .get(&TypeId::of::<HttpMockConfig>())
            .map(|f| f.len())
            .unwrap_or(0);

        if httpmock_configs_len > 1 {
            return Err(TestError::NotSupportedConfig(
                "Only 1 HttpMock can be defined per test".to_string(),
            ));
        }

        let test = Test::current()?;

        let docker = Docker::connect_with_local_defaults()?;
        check_docker(&docker).await?;

        let host = Host::current(&docker).await?;

        Cleanup::new(docker.clone()).purge().await?;

        let mut network = Network::new(docker.clone()).await?;

        if let Some(host_container) = host.container() {
            log::info!("Creating testing environment in containerized mode.");

            // Containerized mode connects the current container into the network.
            network.connect(host_container.id()).await?;
        } else {
            log::info!("Creating testing environment in standalone mode.");
        }

        let starts = self.configs.iter().flat_map(|(_, configs)| {
            configs.values().map(|config| {
                Container::initialized(
                    docker.clone(),
                    test.clone(),
                    host.mode(),
                    &network,
                    config.upcast(),
                )
            })
        });

        let containers = try_join_all(starts)
            .await?
            .into_iter()
            .map(|c| (c.config().hostname().to_string(), c));

        Ok(TestComposite {
            inner: Some(Inner {
                configs: self.configs,
                containers: containers.collect(),
                network,
                test: test.clone(),
            }),
        })
    }
}

impl Drop for TestComposite {
    fn drop(&mut self) {
        let Inner {
            mut network,
            containers,
            test,
            ..
        } = self.inner.take().unwrap();
        tokio::task::block_in_place(|| {
            log::info!("Dropping testing environment.");

            Handle::current().block_on(async {
                if !test.is_success() {
                    tokio::time::sleep(Duration::from_secs(1)).await;
                }
                join_all(containers.into_values().map(|mut container| async move {
                    container.dispose().await;
                }))
                .await;
                network.remove().await;
            })
        });
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use bollard::container::{CreateContainerOptions, NetworkingConfig};
    use bollard::errors::Error as BollardError;
    use bollard::network::CreateNetworkOptions;
    use bollard::secret::EndpointSettings;
    use bollard::Docker;

    use crate::constants::NETWORK_NAME;
    use crate::error::TestError;
    use crate::image::Image;
    use crate::runner::Test;
    use crate::services::httpbin::HttpBinConfig;

    use super::TestComposite;

    #[tokio::test]
    async fn multi_thread_required_error() {
        let result = TestComposite::builder().build().await;
        assert!(matches!(result, Err(TestError::UnavailableMultiThread)));
    }

    #[test]
    fn runtime_required() {
        let result = futures::executor::block_on(TestComposite::builder().build());
        assert!(matches!(result, Err(TestError::UnavailableRuntime(_))));
    }

    #[test]
    fn create_container_logs() {
        let test = Test::builder().module("foo").name("bar").build();

        let target_dir = test.target_dir().to_owned();
        let _ = test.run(async {
            let s1 = HttpBinConfig::builder().hostname("service-1").build();
            let s2 = HttpBinConfig::builder().hostname("service-2").build();

            let _ = TestComposite::builder()
                .with_service(s1)
                .with_service(s2)
                .build()
                .await?;

            assert!(target_dir.join("service-1.log").exists());
            assert!(target_dir.join("service-2.log").exists());

            Ok::<_, TestError>(())
        });

        assert!(!target_dir.join("service-1.log").exists());
        assert!(!target_dir.join("service-2.log").exists());
    }

    #[test]
    fn drop_network() {
        let docker = Docker::connect_with_local_defaults().unwrap();
        let test = Test::builder().module("foo").name("bar").build();

        let _ = test.run(async {
            let s1 = HttpBinConfig::builder().hostname("service-1").build();
            let s2 = HttpBinConfig::builder().hostname("service-2").build();

            let _tc = TestComposite::builder()
                .with_service(s1)
                .with_service(s2)
                .build()
                .await?;

            // Check created network
            let result = docker.inspect_network::<String>(NETWORK_NAME, None).await;
            assert!(result.is_ok());

            Ok::<_, TestError>(())
        });

        // Check network deletion
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let result = runtime.block_on(docker.inspect_network::<String>(NETWORK_NAME, None));

        assert!(matches!(
            result,
            Err(BollardError::DockerResponseServerError {
                status_code: 404,
                ..
            })
        ));
    }

    #[test]
    fn purge_test_assets() -> Result<(), TestError> {
        let test = Test::builder().module("foo").name("bar").build();

        test.run(async {
            let docker = bollard::Docker::connect_with_local_defaults()?;

            // Ensure hello-world image
            let hello_world_image = Image::from_repository("hello-world").with_version("linux");
            hello_world_image.pull(&docker).await?;

            // Create a network that shares the name and is properly labeled.
            let network = docker
                .create_network(CreateNetworkOptions {
                    name: "pdk-test-network",
                    driver: "bridge",
                    labels: HashMap::from([("CreatedBy", "pdk-test")]),
                    ..Default::default()
                })
                .await?;

            let net_id = network.id;

            let hello_world_locator = hello_world_image.locator();
            let hello_world_name = "hello-world";

            // Create a container that uses the network and is properly labeled.
            let container = docker
                .create_container(
                    Some(CreateContainerOptions {
                        name: hello_world_name,
                        platform: None,
                    }),
                    bollard::container::Config {
                        image: Some(hello_world_locator.as_str()),
                        hostname: Some("helloWorld"),
                        network_disabled: Some(false),
                        networking_config: Some(NetworkingConfig {
                            endpoints_config: HashMap::from([(
                                net_id.as_str(),
                                EndpointSettings {
                                    ..Default::default()
                                },
                            )]),
                        }),
                        labels: Some(HashMap::from([("CreatedBy", "pdk-test")])),
                        ..Default::default()
                    },
                )
                .await?;

            // start the container to connect it to the network
            docker.start_container::<&str>(&container.id, None).await?;

            let hello_world_inspect = docker.inspect_container(hello_world_name, None).await;

            // Assert that hello-world container exists
            assert!(hello_world_inspect.is_ok());

            let httpbin_config = HttpBinConfig::builder().hostname("httpbin").build();

            let _composite = TestComposite::builder()
                .with_service(httpbin_config)
                .build()
                .await?;

            let hello_world_inspect = docker.inspect_container(hello_world_name, None).await;

            // Assert that hello-world container no longer exists
            assert!(matches!(
                hello_world_inspect,
                Err(BollardError::DockerResponseServerError {
                    status_code: 404,
                    ..
                })
            ));

            Ok(())
        })
    }
}