Skip to content

Commit 0d14769

Browse files
feat: add with_network() support to ContainerGuard (#226)
- Add with_network() builder method to attach container to a network - Add create_network option (default: true) to auto-create network - Add remove_network_on_drop option to clean up network on drop - Add network() getter to retrieve configured network name - Network is created with bridge driver if it doesn't exist - Network cleanup only happens if guard created the network Closes #219
1 parent f497dd2 commit 0d14769

2 files changed

Lines changed: 176 additions & 3 deletions

File tree

src/testing.rs

Lines changed: 110 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,9 @@
2626
2727
use crate::command::DockerCommand;
2828
use crate::template::{Template, TemplateError};
29-
use crate::{LogsCommand, PortCommand, RmCommand, StopCommand};
29+
use crate::{
30+
LogsCommand, NetworkCreateCommand, NetworkRmCommand, PortCommand, RmCommand, StopCommand,
31+
};
3032
use std::sync::atomic::{AtomicBool, Ordering};
3133
use std::sync::Arc;
3234

@@ -46,6 +48,12 @@ pub struct GuardOptions {
4648
pub reuse_if_running: bool,
4749
/// Automatically wait for container to be ready after start (default: false)
4850
pub wait_for_ready: bool,
51+
/// Network to attach the container to (default: None)
52+
pub network: Option<String>,
53+
/// Create the network if it doesn't exist (default: true when network is set)
54+
pub create_network: bool,
55+
/// Remove the network on drop (default: false)
56+
pub remove_network_on_drop: bool,
4957
}
5058

5159
impl Default for GuardOptions {
@@ -57,6 +65,9 @@ impl Default for GuardOptions {
5765
capture_logs: false,
5866
reuse_if_running: false,
5967
wait_for_ready: false,
68+
network: None,
69+
create_network: true,
70+
remove_network_on_drop: false,
6071
}
6172
}
6273
}
@@ -147,6 +158,53 @@ impl<T: Template> ContainerGuardBuilder<T> {
147158
self
148159
}
149160

161+
/// Attach the container to a Docker network.
162+
///
163+
/// By default, the network will be created if it doesn't exist. Use
164+
/// `create_network(false)` to disable automatic network creation.
165+
///
166+
/// # Example
167+
///
168+
/// ```rust,no_run
169+
/// # use docker_wrapper::testing::ContainerGuard;
170+
/// # use docker_wrapper::RedisTemplate;
171+
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
172+
/// let guard = ContainerGuard::new(RedisTemplate::new("redis"))
173+
/// .with_network("test-network")
174+
/// .start()
175+
/// .await?;
176+
/// // Container is attached to "test-network"
177+
/// # Ok(())
178+
/// # }
179+
/// ```
180+
#[must_use]
181+
pub fn with_network(mut self, network: impl Into<String>) -> Self {
182+
self.options.network = Some(network.into());
183+
self
184+
}
185+
186+
/// Set whether to create the network if it doesn't exist (default: true).
187+
///
188+
/// Only applies when a network is specified via `with_network()`.
189+
#[must_use]
190+
pub fn create_network(mut self, create: bool) -> Self {
191+
self.options.create_network = create;
192+
self
193+
}
194+
195+
/// Set whether to remove the network on drop (default: false).
196+
///
197+
/// This is useful for cleaning up test-specific networks. Only applies
198+
/// when a network is specified via `with_network()`.
199+
///
200+
/// Note: The network removal will fail silently if other containers are
201+
/// still using it.
202+
#[must_use]
203+
pub fn remove_network_on_drop(mut self, remove: bool) -> Self {
204+
self.options.remove_network_on_drop = remove;
205+
self
206+
}
207+
150208
/// Start the container and return a guard that manages its lifecycle.
151209
///
152210
/// If `reuse_if_running` is enabled and a container is already running,
@@ -155,11 +213,33 @@ impl<T: Template> ContainerGuardBuilder<T> {
155213
/// If `wait_for_ready` is enabled, this method will block until the
156214
/// container passes its readiness check.
157215
///
216+
/// If a network is specified via `with_network()`, the container will be
217+
/// attached to that network. The network will be created if it doesn't
218+
/// exist (unless `create_network(false)` was called).
219+
///
158220
/// # Errors
159221
///
160222
/// Returns an error if the container fails to start or the readiness check times out.
161-
pub async fn start(self) -> Result<ContainerGuard<T>, TemplateError> {
223+
pub async fn start(mut self) -> Result<ContainerGuard<T>, TemplateError> {
162224
let wait_for_ready = self.options.wait_for_ready;
225+
let mut network_created = false;
226+
227+
// Create network if specified and create_network is enabled
228+
if let Some(ref network) = self.options.network {
229+
if self.options.create_network {
230+
// Try to create the network (ignore errors if it already exists)
231+
let result = NetworkCreateCommand::new(network)
232+
.driver("bridge")
233+
.execute()
234+
.await;
235+
236+
// Track if we successfully created it (for cleanup purposes)
237+
network_created = result.is_ok();
238+
}
239+
240+
// Set the network on the template
241+
self.template.config_mut().network = Some(network.clone());
242+
}
163243

164244
// Check if we should reuse an existing container
165245
if self.options.reuse_if_running {
@@ -169,6 +249,7 @@ impl<T: Template> ContainerGuardBuilder<T> {
169249
container_id: None, // We don't have the ID for reused containers
170250
options: self.options,
171251
was_reused: true,
252+
network_created,
172253
cleaned_up: Arc::new(AtomicBool::new(false)),
173254
};
174255

@@ -189,6 +270,7 @@ impl<T: Template> ContainerGuardBuilder<T> {
189270
container_id: Some(container_id),
190271
options: self.options,
191272
was_reused: false,
273+
network_created,
192274
cleaned_up: Arc::new(AtomicBool::new(false)),
193275
};
194276

@@ -229,6 +311,7 @@ pub struct ContainerGuard<T: Template> {
229311
container_id: Option<String>,
230312
options: GuardOptions,
231313
was_reused: bool,
314+
network_created: bool,
232315
cleaned_up: Arc<AtomicBool>,
233316
}
234317

@@ -262,6 +345,12 @@ impl<T: Template> ContainerGuard<T> {
262345
self.was_reused
263346
}
264347

348+
/// Get the network name, if one was configured.
349+
#[must_use]
350+
pub fn network(&self) -> Option<&str> {
351+
self.options.network.as_deref()
352+
}
353+
265354
/// Check if the container is currently running.
266355
///
267356
/// # Errors
@@ -437,9 +526,11 @@ impl<T: Template> Drop for ContainerGuard<T> {
437526
// Perform cleanup - need to spawn a runtime since Drop isn't async
438527
let should_stop = self.options.stop_on_drop;
439528
let should_remove = self.options.remove_on_drop;
529+
let should_remove_network = self.options.remove_network_on_drop && self.network_created;
440530
let container_name = self.template.config().name.clone();
531+
let network_name = self.options.network.clone();
441532

442-
if !should_stop && !should_remove {
533+
if !should_stop && !should_remove && !should_remove_network {
443534
return;
444535
}
445536

@@ -448,6 +539,7 @@ impl<T: Template> Drop for ContainerGuard<T> {
448539
if tokio::runtime::Handle::try_current().is_ok() {
449540
// We're in an async context - use spawn_blocking to avoid blocking the runtime
450541
let container_name_clone = container_name.clone();
542+
let network_name_clone = network_name.clone();
451543
let _ = std::thread::spawn(move || {
452544
let rt = tokio::runtime::Builder::new_current_thread()
453545
.enable_all()
@@ -460,6 +552,12 @@ impl<T: Template> Drop for ContainerGuard<T> {
460552
if should_remove {
461553
let _ = RmCommand::new(&container_name_clone).force().run().await;
462554
}
555+
// Remove network after container (network must be empty)
556+
if should_remove_network {
557+
if let Some(ref network) = network_name_clone {
558+
let _ = NetworkRmCommand::new(network).execute().await;
559+
}
560+
}
463561
});
464562
})
465563
.join();
@@ -476,6 +574,12 @@ impl<T: Template> Drop for ContainerGuard<T> {
476574
if should_remove {
477575
let _ = RmCommand::new(&container_name).force().run().await;
478576
}
577+
// Remove network after container (network must be empty)
578+
if should_remove_network {
579+
if let Some(ref network) = network_name {
580+
let _ = NetworkRmCommand::new(network).execute().await;
581+
}
582+
}
479583
});
480584
}
481585
}
@@ -495,6 +599,9 @@ mod tests {
495599
assert!(!opts.capture_logs);
496600
assert!(!opts.reuse_if_running);
497601
assert!(!opts.wait_for_ready);
602+
assert!(opts.network.is_none());
603+
assert!(opts.create_network);
604+
assert!(!opts.remove_network_on_drop);
498605
}
499606

500607
#[test]

tests/container_guard_integration.rs

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,3 +217,69 @@ async fn test_container_guard_auto_wait_for_ready() {
217217

218218
// Container cleanup happens on drop
219219
}
220+
221+
#[tokio::test]
222+
async fn test_container_guard_with_network() {
223+
let name = unique_name("guard-network");
224+
let network_name = unique_name("test-network");
225+
226+
let guard = ContainerGuard::new(RedisTemplate::new(&name).port(next_port()))
227+
.with_network(&network_name)
228+
.remove_network_on_drop(true)
229+
.start()
230+
.await
231+
.expect("Failed to start container");
232+
233+
// Container should be running
234+
assert!(
235+
guard.is_running().await.expect("Failed to check running"),
236+
"Container should be running"
237+
);
238+
239+
// Network should be set
240+
assert_eq!(guard.network(), Some(network_name.as_str()));
241+
242+
// Verify the container is on the network by checking the template config
243+
assert_eq!(
244+
guard.template().config().network,
245+
Some(network_name.clone())
246+
);
247+
248+
// Container and network cleanup happens on drop
249+
}
250+
251+
#[tokio::test]
252+
async fn test_container_guard_network_no_auto_create() {
253+
use docker_wrapper::{DockerCommand, NetworkCreateCommand, NetworkRmCommand};
254+
255+
let name = unique_name("guard-network-manual");
256+
let network_name = unique_name("manual-network");
257+
258+
// Create the network manually first
259+
NetworkCreateCommand::new(&network_name)
260+
.driver("bridge")
261+
.execute()
262+
.await
263+
.expect("Failed to create network");
264+
265+
let guard = ContainerGuard::new(RedisTemplate::new(&name).port(next_port()))
266+
.with_network(&network_name)
267+
.create_network(false) // Don't auto-create
268+
.start()
269+
.await
270+
.expect("Failed to start container");
271+
272+
assert!(
273+
guard.is_running().await.expect("Failed to check running"),
274+
"Container should be running"
275+
);
276+
277+
// Drop the guard (container will be removed but network won't since we didn't create it)
278+
drop(guard);
279+
280+
// Give Docker a moment to clean up
281+
tokio::time::sleep(Duration::from_millis(500)).await;
282+
283+
// Clean up network manually
284+
let _ = NetworkRmCommand::new(&network_name).execute().await;
285+
}

0 commit comments

Comments
 (0)