From e97fce685af31f6fbd16ad662f80f8d2ca9de878 Mon Sep 17 00:00:00 2001 From: Mark S Date: Wed, 20 Nov 2024 09:01:39 -0500 Subject: [PATCH] feat: add support for user-configured labels --- testcontainers/src/core/containers/request.rs | 7 +++ testcontainers/src/core/image/image_ext.rs | 27 ++++++++ testcontainers/src/runners/async_runner.rs | 61 ++++++++++++++++--- 3 files changed, 88 insertions(+), 7 deletions(-) diff --git a/testcontainers/src/core/containers/request.rs b/testcontainers/src/core/containers/request.rs index 81ceffaa..24595565 100644 --- a/testcontainers/src/core/containers/request.rs +++ b/testcontainers/src/core/containers/request.rs @@ -25,6 +25,7 @@ pub struct ContainerRequest { pub(crate) image_tag: Option, pub(crate) container_name: Option, pub(crate) network: Option, + pub(crate) labels: BTreeMap, pub(crate) env_vars: BTreeMap, pub(crate) hosts: BTreeMap, pub(crate) mounts: Vec, @@ -74,6 +75,10 @@ impl ContainerRequest { &self.network } + pub fn labels(&self) -> &BTreeMap { + &self.labels + } + pub fn container_name(&self) -> &Option { &self.container_name } @@ -190,6 +195,7 @@ impl From for ContainerRequest { image_tag: None, container_name: None, network: None, + labels: BTreeMap::default(), env_vars: BTreeMap::default(), hosts: BTreeMap::default(), mounts: Vec::new(), @@ -235,6 +241,7 @@ impl Debug for ContainerRequest { .field("image_tag", &self.image_tag) .field("container_name", &self.container_name) .field("network", &self.network) + .field("labels", &self.labels) .field("env_vars", &self.env_vars) .field("hosts", &self.hosts) .field("mounts", &self.mounts) diff --git a/testcontainers/src/core/image/image_ext.rs b/testcontainers/src/core/image/image_ext.rs index bdbd12af..0bff2831 100644 --- a/testcontainers/src/core/image/image_ext.rs +++ b/testcontainers/src/core/image/image_ext.rs @@ -48,6 +48,18 @@ pub trait ImageExt { /// Sets the network the container will be connected to. fn with_network(self, network: impl Into) -> ContainerRequest; + /// Adds the specified labels to the container. + /// + /// **Note**: in addition to all keys in the `com.testcontainers.*` namespace, there + /// are certain labels that are used by `testcontainers` internally which will always + /// be unconditionally overwritten, and so should not be expected or relied upon to + /// be applied correctly if they are included in `labels`. Currently, they are: + /// - `managed-by` + fn with_labels( + self, + labels: impl IntoIterator, impl Into)>, + ) -> ContainerRequest; + /// Adds an environment variable to the container. fn with_env_var(self, name: impl Into, value: impl Into) -> ContainerRequest; @@ -164,6 +176,21 @@ impl>, I: Image> ImageExt for RI { } } + fn with_labels( + self, + labels: impl IntoIterator, impl Into)>, + ) -> ContainerRequest { + let mut container_req = self.into(); + + container_req.labels.extend( + labels + .into_iter() + .map(|(key, value)| (key.into(), value.into())), + ); + + container_req + } + fn with_env_var( self, name: impl Into, diff --git a/testcontainers/src/runners/async_runner.rs b/testcontainers/src/runners/async_runner.rs index 47e9942e..d0b36d6a 100644 --- a/testcontainers/src/runners/async_runner.rs +++ b/testcontainers/src/runners/async_runner.rs @@ -1,12 +1,5 @@ use std::{collections::HashMap, time::Duration}; -use async_trait::async_trait; -use bollard::{ - container::{Config, CreateContainerOptions}, - models::{HostConfig, PortBinding}, -}; -use bollard_stubs::models::{HostConfigCgroupnsModeEnum, ResourcesUlimits}; - use crate::{ core::{ client::{Client, ClientError}, @@ -18,6 +11,12 @@ use crate::{ }, ContainerAsync, ContainerRequest, Image, }; +use async_trait::async_trait; +use bollard::{ + container::{Config, CreateContainerOptions}, + models::{HostConfig, PortBinding}, +}; +use bollard_stubs::models::{HostConfigCgroupnsModeEnum, ResourcesUlimits}; const DEFAULT_STARTUP_TIMEOUT: Duration = Duration::from_secs(60); @@ -63,8 +62,17 @@ where .map(|(key, value)| format!("{key}:{value}")) .collect(); + let labels = HashMap::::from_iter( + container_req + .labels() + .iter() + .map(|(key, value)| (key.into(), value.into())) + .chain([("managed-by".to_string(), "testcontainers".to_string())]), + ); + let mut config: Config = Config { image: Some(container_req.descriptor()), + labels: Some(labels), host_config: Some(HostConfig { privileged: Some(container_req.privileged()), extra_hosts: Some(extra_hosts), @@ -297,6 +305,45 @@ mod tests { ImageExt, }; + /// Test that all user-supplied labels are added to containers started by `AsyncRunner::start` + #[tokio::test] + async fn async_start_should_apply_expected_labels() -> anyhow::Result<()> { + let mut labels = HashMap::from([ + ("foo".to_string(), "bar".to_string()), + ("baz".to_string(), "qux".to_string()), + ("managed-by".to_string(), "the-time-wizard".to_string()), + ]); + + let container = GenericImage::new("hello-world", "latest") + .with_labels(&labels) + .start() + .await?; + + let client = Client::lazy_client().await?; + + let container_labels = client + .inspect(container.id()) + .await? + .config + .unwrap_or_default() + .labels + .unwrap_or_default(); + + // the created labels and container labels shouldn't actually be identical, + // as the `managed-by: testcontainers` label is always unconditionally applied + // to all containers by `AsyncRunner::start`, with the value `testcontainers` + // being applied *last* explicitly so that even user-supplied values of + // the `managed-by` key will be overwritten + assert_ne!(&labels, &container_labels); + + // If we add the expected `managed-by` value though, they should then match + labels.insert("managed-by".to_string(), "testcontainers".to_string()); + + assert_eq!(labels, container_labels); + + Ok(()) + } + #[tokio::test] async fn async_run_command_should_expose_all_ports_if_no_explicit_mapping_requested( ) -> anyhow::Result<()> {