From 4bc6777f19c52750a7c544b79a457e5c81276bdb Mon Sep 17 00:00:00 2001
From: Dale Hui <dhui@users.noreply.github.com>
Date: Sun, 14 Apr 2024 23:13:36 -0700
Subject: [PATCH] Add dktesting.Cleanup() method

---
 dktesting/dktesting.go | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/dktesting/dktesting.go b/dktesting/dktesting.go
index 342af2fdc..7a0785860 100644
--- a/dktesting/dktesting.go
+++ b/dktesting/dktesting.go
@@ -1,11 +1,15 @@
 package dktesting
 
 import (
+	"context"
+	"fmt"
 	"testing"
 )
 
 import (
 	"github.com/dhui/dktest"
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/client"
 )
 
 // ContainerSpec holds Docker testing setup specifications
@@ -14,6 +18,26 @@ type ContainerSpec struct {
 	Options   dktest.Options
 }
 
+// Cleanup cleanups the ContainerSpec after a test run by removing the ContainerSpec's image
+func (s *ContainerSpec) Cleanup() (retErr error) {
+	// copied from dktest.RunContext()
+	dc, err := client.NewClientWithOpts(client.FromEnv, client.WithVersion("1.41"))
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := dc.Close(); err != nil && retErr == nil {
+			retErr = fmt.Errorf("error closing Docker client: %w", err)
+		}
+	}()
+	ctx, timeoutCancelFunc := context.WithTimeout(context.Background(), s.Options.CleanupTimeout)
+	defer timeoutCancelFunc()
+	if _, err := dc.ImageRemove(ctx, s.ImageName, types.ImageRemoveOptions{Force: true, PruneChildren: true}); err != nil {
+		return err
+	}
+	return nil
+}
+
 // ParallelTest runs Docker tests in parallel
 func ParallelTest(t *testing.T, specs []ContainerSpec,
 	testFunc func(*testing.T, dktest.ContainerInfo)) {