pdk-test 1.6.0

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

//! Network management for integration testing
//!
//! This module provides utilities for managing Docker networks in integration tests.
//! It handles network creation, connection and removal with proper cleanup.
//!
//! ## Primary types
//!
//! - [`Network`]: manages Docker network lifecycle
//! - [`NetworkStatus`]: enum for network status
//!

use std::collections::{HashMap, HashSet};

use crate::constants::{LABEL_KEY, LABEL_VALUE, NETWORK_NAME};
use bollard::errors::Error as DockerError;
use bollard::network::{ConnectNetworkOptions, CreateNetworkOptions, DisconnectNetworkOptions};
use bollard::service::EndpointSettings;
use bollard::Docker;
use futures::future::try_join_all;
use tokio::runtime::Handle;

use crate::error::TestError;

enum NetworkStatus {
    Created,
    Removed,
}

/// Struct to manage the network.
pub struct Network {
    id: String,
    docker: Docker,
    status: NetworkStatus,
    connected_containers: HashSet<String>,
}

impl Network {
    /// Create a new network.
    pub async fn new(docker: Docker) -> Result<Self, TestError> {
        log::info!("Creating Docker network.");

        let response = docker
            .create_network(CreateNetworkOptions {
                name: NETWORK_NAME,
                driver: "bridge",
                labels: HashMap::from([(LABEL_KEY, LABEL_VALUE)]),
                ..Default::default()
            })
            .await?;
        Ok(Self {
            id: response.id,
            docker,
            status: NetworkStatus::Created,
            connected_containers: Default::default(),
        })
    }

    /// Get the ID of the network.
    pub fn id(&self) -> &str {
        &self.id
    }

    /// Connect a container to the network.
    pub async fn connect(&mut self, container_id: &str) -> Result<(), TestError> {
        self.docker
            .connect_network(
                &self.id,
                ConnectNetworkOptions {
                    container: container_id,
                    endpoint_config: EndpointSettings {
                        ..Default::default()
                    },
                },
            )
            .await?;

        self.connected_containers.insert(container_id.to_string());

        Ok(())
    }

    async fn try_remove(&mut self) -> Result<(), TestError> {
        if matches!(self.status, NetworkStatus::Removed) {
            return Ok(());
        }

        let network_id = &self.id;
        let docker = &self.docker;

        let disconnect = self.connected_containers.iter().map(|c| async move {
            let result = docker
                .disconnect_network(
                    network_id,
                    DisconnectNetworkOptions {
                        container: c,
                        force: true,
                    },
                )
                .await;

            match result {
                // Ignore already deleted containers
                Err(DockerError::DockerResponseServerError {
                    status_code: 404, ..
                }) => Ok(()),
                r => r,
            }
        });
        try_join_all(disconnect).await?;

        self.docker.remove_network(&self.id).await?;

        self.status = NetworkStatus::Removed;

        Ok(())
    }

    /// Remove the network.
    pub async fn remove(&mut self) {
        log::info!("Removing Docker network.");

        if let Err(e) = self.try_remove().await {
            log::error!("unable to remove network with id {}: {e}", self.id);
        }
    }
}

impl Drop for Network {
    fn drop(&mut self) {
        tokio::task::block_in_place(|| Handle::current().block_on(self.remove()));
    }
}