From fcccd5c36d1ff2ec018686723aa1d3f0ee02a552 Mon Sep 17 00:00:00 2001
From: Dale Hui <dhui@users.noreply.github.com>
Date: Sun, 14 Apr 2024 22:19:03 -0700
Subject: [PATCH] Make image removal optional (defaults to false)

---
 dktest.go  |  4 +++-
 options.go | 11 +++++++----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/dktest.go b/dktest.go
index c4564cc..e30e4e9 100644
--- a/dktest.go
+++ b/dktest.go
@@ -217,7 +217,9 @@ func RunContext(ctx context.Context, logger Logger, imgName string, opts Options
 			stopCtx, stopTimeoutCancelFunc := context.WithTimeout(ctx, opts.CleanupTimeout)
 			defer stopTimeoutCancelFunc()
 			stopContainer(stopCtx, logger, dc, c, opts.LogStdout, opts.LogStderr)
-			removeImage(ctx, logger, dc, imgName)
+			if opts.CleanupImage {
+				removeImage(stopCtx, logger, dc, imgName)
+			}
 		}()
 
 		if waitContainerReady(runCtx, logger, c, opts.ReadyFunc, opts.ReadyTimeout) {
diff --git a/options.go b/options.go
index 05fb1bd..52c5ba3 100644
--- a/options.go
+++ b/options.go
@@ -19,10 +19,13 @@ type Options struct {
 	ReadyTimeout time.Duration
 	// CleanupTimeout is the timeout used when stopping and removing a container
 	CleanupTimeout time.Duration
-	ReadyFunc      func(context.Context, ContainerInfo) bool
-	Env            map[string]string
-	Entrypoint     []string
-	Cmd            []string
+	// CleanupImage specifies whether or not the image should be removed after the test run.
+	// If the image is used by multiple tests, you'll want to cleanup the image yourself.
+	CleanupImage bool
+	ReadyFunc    func(context.Context, ContainerInfo) bool
+	Env          map[string]string
+	Entrypoint   []string
+	Cmd          []string
 	// If you prefer to specify your port bindings as a string, use nat.ParsePortSpecs()
 	PortBindings nat.PortMap
 	PortRequired bool