Skip to content

Commit

Permalink
feat: add support for user-configured labels
Browse files Browse the repository at this point in the history
  • Loading branch information
the-wondersmith committed Nov 20, 2024
1 parent 2f2dca1 commit e97fce6
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 7 deletions.
7 changes: 7 additions & 0 deletions testcontainers/src/core/containers/request.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ pub struct ContainerRequest<I: Image> {
pub(crate) image_tag: Option<String>,
pub(crate) container_name: Option<String>,
pub(crate) network: Option<String>,
pub(crate) labels: BTreeMap<String, String>,
pub(crate) env_vars: BTreeMap<String, String>,
pub(crate) hosts: BTreeMap<String, Host>,
pub(crate) mounts: Vec<Mount>,
Expand Down Expand Up @@ -74,6 +75,10 @@ impl<I: Image> ContainerRequest<I> {
&self.network
}

pub fn labels(&self) -> &BTreeMap<String, String> {
&self.labels
}

pub fn container_name(&self) -> &Option<String> {
&self.container_name
}
Expand Down Expand Up @@ -190,6 +195,7 @@ impl<I: Image> From<I> for ContainerRequest<I> {
image_tag: None,
container_name: None,
network: None,
labels: BTreeMap::default(),
env_vars: BTreeMap::default(),
hosts: BTreeMap::default(),
mounts: Vec::new(),
Expand Down Expand Up @@ -235,6 +241,7 @@ impl<I: Image + Debug> Debug for ContainerRequest<I> {
.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)
Expand Down
27 changes: 27 additions & 0 deletions testcontainers/src/core/image/image_ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,18 @@ pub trait ImageExt<I: Image> {
/// Sets the network the container will be connected to.
fn with_network(self, network: impl Into<String>) -> ContainerRequest<I>;

/// 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<Item = (impl Into<String>, impl Into<String>)>,
) -> ContainerRequest<I>;

/// Adds an environment variable to the container.
fn with_env_var(self, name: impl Into<String>, value: impl Into<String>)
-> ContainerRequest<I>;
Expand Down Expand Up @@ -164,6 +176,21 @@ impl<RI: Into<ContainerRequest<I>>, I: Image> ImageExt<I> for RI {
}
}

fn with_labels(
self,
labels: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
) -> ContainerRequest<I> {
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<String>,
Expand Down
61 changes: 54 additions & 7 deletions testcontainers/src/runners/async_runner.rs
Original file line number Diff line number Diff line change
@@ -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},
Expand All @@ -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);

Expand Down Expand Up @@ -63,8 +62,17 @@ where
.map(|(key, value)| format!("{key}:{value}"))
.collect();

let labels = HashMap::<String, String>::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<String> = Config {
image: Some(container_req.descriptor()),
labels: Some(labels),
host_config: Some(HostConfig {
privileged: Some(container_req.privileged()),
extra_hosts: Some(extra_hosts),
Expand Down Expand Up @@ -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<()> {
Expand Down

0 comments on commit e97fce6

Please sign in to comment.