pdk-test 1.7.0

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

//! Image management for integration testing
//!
//! This module provides utilities for managing Docker images in integration tests.
//! It handles image lookup, pulling and creation with proper error handling.
//!

use crate::constants::{DOCKER_PLATFORM, LATEST_VERSION};
use bollard::errors::Error as BollardError;
use bollard::{auth::DockerCredentials, image::CreateImageOptions, Docker};
use futures::{StreamExt as _, TryStreamExt as _};

use crate::error::TestError;

const DOCKER_HUB: &str = "https://index.docker.io/v1/";

/// Struct to manage Docker image lifecycle.
#[derive(Clone, Debug)]
pub struct Image {
    repository: String,
    version: String,
    server: Option<String>,
}

impl Image {
    /// Create a new image from a repository.
    pub fn from_repository<T: Into<String>>(repository: T) -> Self {
        Self {
            repository: repository.into(),
            version: LATEST_VERSION.into(),
            server: None,
        }
    }

    /// Set the version of the image.
    pub fn with_version<T: Into<String>>(self, version: T) -> Self {
        Self {
            version: version.into(),
            ..self
        }
    }

    /// Get the repository of the image.
    pub fn repository(&self) -> &str {
        self.repository.as_str()
    }

    /// Get the version of the image.
    pub fn version(&self) -> &str {
        self.version.as_str()
    }

    /// Get the locator of the image.
    pub fn locator(&self) -> String {
        format!("{}:{}", self.repository, self.version)
    }

    /// Pull the image from the repository.
    pub(super) async fn pull(&self, docker: &Docker) -> Result<(), TestError> {
        log::info!("Looking up for Docker image from {}.", self.repository);

        let locator = self.locator();

        // Always try to inspect image locally first to avoid unnecessary pulls
        let inspect = docker.inspect_image(&locator).await;
        match inspect {
            // Image already exists locally
            Ok(_) => {
                log::info!("Image {locator} found locally, skipping pull.");
                return Ok(());
            }

            // Must pull image
            Err(BollardError::DockerResponseServerError {
                status_code: 404, ..
            }) => {
                if self.version() == LATEST_VERSION {
                    log::warn!(
                        "Avoid using '{LATEST_VERSION}' version for Docker image {locator}. \
                                Using exact versions could improve test performance."
                    );
                }
            }

            // Unmanaged error
            Err(e) => {
                return Err(e.into());
            }
        }

        let server = self.server.as_deref().unwrap_or(DOCKER_HUB);
        let creds = match docker_credential::get_credential(server) {
            Ok(creds) => Some(creds),
            Err(err) => {
                log::warn!("Unable to find Docker credentials: {err}.");
                None
            }
        };
        let creds = creds.map(|creds| match creds {
            docker_credential::DockerCredential::IdentityToken(token) => DockerCredentials {
                identitytoken: Some(token),
                ..Default::default()
            },
            docker_credential::DockerCredential::UsernamePassword(username, password) => {
                DockerCredentials {
                    username: Some(username),
                    password: Some(password),
                    serveraddress: Some(server.to_string()),
                    ..Default::default()
                }
            }
        });

        let platform = std::env::var(DOCKER_PLATFORM).unwrap_or_default();

        log::info!("Pulling Docker image from {}.", self.repository);

        let _ = docker
            .create_image(
                Some(CreateImageOptions {
                    from_image: locator.as_str(),
                    platform: platform.as_str(),
                    ..Default::default()
                }),
                None,
                creds,
            )
            .map(|r| match r {
                Ok(_) => Ok(()),
                // Ignore error due to a known Bollard issue https://github.com/fussybeaver/bollard/pull/427
                Err(err @ BollardError::JsonSerdeError { .. }) => {
                    log::warn!(
                        "Unable to deserialize Docker API result when pulling {locator}: {err}."
                    );
                    Ok(())
                }
                Err(e) => Err(e),
            })
            .try_collect::<Vec<_>>()
            .await?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::future::Future;

    use bollard::{secret::ImageInspect, Docker, API_DEFAULT_VERSION};
    use httpmock::MockServer;
    use serde_json::json;

    use crate::TestError;

    use super::Image;

    // Runs async test with mocked docker daemon
    fn run_test<T, F>(body: T)
    where
        T: FnOnce(MockServer, Docker) -> F,
        F: Future<Output = ()>,
    {
        tokio::runtime::Runtime::new().unwrap().block_on(async {
            let docker_mock = MockServer::start_async().await;

            let docker =
                Docker::connect_with_http(&docker_mock.base_url(), 30, API_DEFAULT_VERSION)
                    .unwrap();

            body(docker_mock, docker).await;
        })
    }

    #[test]
    fn skip_create_image() {
        run_test(|docker_mock, docker| async move {
            // mock inspect
            let inspect_mock = docker_mock
                .mock_async(|when, then| {
                    when.path_contains("/images/helloworld:1.0/json")
                        .method(httpmock::Method::GET);
                    then.status(200).json_body_obj(&ImageInspect {
                        ..Default::default()
                    });
                })
                .await;

            // mock create
            let create_mock = docker_mock
                .mock_async(|when, _then| {
                    when.path_contains("/images/create")
                        .method(httpmock::Method::POST);
                })
                .await;

            // Pull image
            Image::from_repository("helloworld")
                .with_version("1.0")
                .pull(&docker)
                .await
                .unwrap();

            // Assertions
            inspect_mock.assert_hits_async(1).await;
            create_mock.assert_hits_async(0).await;
        })
    }

    #[test]
    fn pull_image_when_not_found() {
        run_test(|docker_mock, docker| async move {
            // mock inspect
            let inspect_mock = docker_mock
                .mock_async(|when, then| {
                    when.path_contains("/images/helloworld:1.0/json")
                        .method(httpmock::Method::GET);

                    // Image not found
                    then.status(404).json_body(json!({
                        "message": "not found"
                    }));
                })
                .await;

            // mock create
            let create_mock = docker_mock
                .mock_async(|when, then| {
                    when.path_contains("/images/create")
                        .method(httpmock::Method::POST);

                    // This mock returns non-JSON payload because tested Docker API versions
                    // do not return valid JSON when creating images.
                    then.status(200).body("Image created");
                })
                .await;

            // Pull image
            Image::from_repository("helloworld")
                .with_version("1.0")
                .pull(&docker)
                .await
                .unwrap();

            // Assertions
            inspect_mock.assert_hits_async(1).await;
            create_mock.assert_hits_async(1).await;
        })
    }

    #[test]
    fn inspect_image_before_pull() {
        run_test(|docker_mock, docker| async move {
            // mock inspect
            let inspect_mock = docker_mock
                .mock_async(|when, then| {
                    when.path_contains("/images/helloworld:latest/json")
                        .method(httpmock::Method::GET);
                    // Image not found locally, so it will proceed to pull
                    then.status(404).json_body(json!({
                        "message": "not found"
                    }));
                })
                .await;

            // mock create
            let create_mock = docker_mock
                .mock_async(|when, then| {
                    when.path_contains("/images/create")
                        .method(httpmock::Method::POST);

                    // This mock returns non-JSON payload because tested Docker API versions
                    // do not return valid JSON when creating images.
                    then.status(200).body("Image created");
                })
                .await;

            // Pull image
            Image::from_repository("helloworld")
                .with_version("latest")
                .pull(&docker)
                .await
                .unwrap();

            // Assertions
            inspect_mock.assert_hits_async(1).await; // Always called
            create_mock.assert_hits_async(1).await;
        })
    }

    #[test]
    fn inspect_error_bubbling() {
        run_test(|docker_mock, docker| async move {
            // mock inspect
            let inspect_mock = docker_mock
                .mock_async(|when, then| {
                    when.path_contains("/images/helloworld:1.0/json")
                        .method(httpmock::Method::GET);

                    then.status(405).json_body(json!({
                        "message": "unknown problem"
                    }));
                })
                .await;

            // Pull image
            let result = Image::from_repository("helloworld")
                .with_version("1.0")
                .pull(&docker)
                .await;

            assert!(matches!(result, Err(TestError::Docker(_))));
            inspect_mock.assert_hits_async(1).await;
        })
    }
}