From 78b44f2115a01936f56f1f77105235b8c3de0054 Mon Sep 17 00:00:00 2001
From: Josh Wolf <josh@wolfs.io>
Date: Tue, 9 Apr 2024 16:55:51 -0400
Subject: [PATCH] plumb through WithTestRemove option

Signed-off-by: Josh Wolf <josh@wolfs.io>
---
 pkg/build/test.go         | 28 +++++++++++++++++++++++-----
 pkg/build/test_options.go |  9 +++++++++
 pkg/cli/test.go           |  2 +-
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/pkg/build/test.go b/pkg/build/test.go
index 5587b213e..851e1ae80 100644
--- a/pkg/build/test.go
+++ b/pkg/build/test.go
@@ -62,13 +62,16 @@ type Test struct {
 	DebugRunner       bool
 	Interactive       bool
 	LogPolicy         []string
+	Remove            bool
+	containerConfigs  []*container.Config
 }
 
 func NewTest(ctx context.Context, opts ...TestOption) (*Test, error) {
 	t := Test{
-		WorkspaceIgnore: ".melangeignore",
-		Arch:            apko_types.ParseArchitecture(runtime.GOARCH),
-		LogPolicy:       []string{"builtin:stderr"},
+		WorkspaceIgnore:  ".melangeignore",
+		Arch:             apko_types.ParseArchitecture(runtime.GOARCH),
+		LogPolicy:        []string{"builtin:stderr"},
+		containerConfigs: []*container.Config{},
 	}
 
 	for _, opt := range opts {
@@ -117,8 +120,21 @@ func NewTest(ctx context.Context, opts ...TestOption) (*Test, error) {
 	return &t, nil
 }
 
-func (t *Test) Close() error {
-	return t.Runner.Close()
+func (t *Test) Close(ctx context.Context) error {
+	log := clog.FromContext(ctx)
+	errs := []error{}
+	if t.Remove {
+		log.Infof("deleting guest dir %s", t.GuestDir)
+		errs = append(errs, os.RemoveAll(t.GuestDir))
+		log.Infof("deleting workspace dir %s", t.WorkspaceDir)
+		errs = append(errs, os.RemoveAll(t.WorkspaceDir))
+		for _, cfg := range t.containerConfigs {
+			errs = append(errs, t.Runner.OCIImageLoader().RemoveImage(context.WithoutCancel(ctx), cfg.ImgRef))
+		}
+	}
+	errs = append(errs, t.Runner.Close())
+
+	return errors.Join(errs...)
 }
 
 // BuildGuest invokes apko to create the test image for the guest environment.
@@ -423,6 +439,7 @@ func (t *Test) TestPackage(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("unable to build workspace config: %w", err)
 	}
+	t.containerConfigs = append(t.containerConfigs, cfg)
 
 	if !t.IsTestless() {
 		cfg.Arch = t.Arch
@@ -489,6 +506,7 @@ func (t *Test) TestPackage(ctx context.Context) error {
 			if err != nil {
 				return fmt.Errorf("unable to build workspace config: %w", err)
 			}
+			t.containerConfigs = append(t.containerConfigs, subCfg)
 			subCfg.Arch = t.Arch
 			if err := t.Runner.StartPod(ctx, subCfg); err != nil {
 				return fmt.Errorf("unable to start subpackage test pod: %w", err)
diff --git a/pkg/build/test_options.go b/pkg/build/test_options.go
index 06525e9f5..15c1bbcd6 100644
--- a/pkg/build/test_options.go
+++ b/pkg/build/test_options.go
@@ -182,3 +182,12 @@ func WithExtraTestPackages(extraTestPackages []string) TestOption {
 		return nil
 	}
 }
+
+// WithRemove indicates whether the the test will clean up after itself.
+// This includes deleting any intermediate artifacts like container images.
+func WithTestRemove(remove bool) TestOption {
+	return func(b *Test) error {
+		b.Remove = remove
+		return nil
+	}
+}
diff --git a/pkg/cli/test.go b/pkg/cli/test.go
index 1028b9497..770084b05 100644
--- a/pkg/cli/test.go
+++ b/pkg/cli/test.go
@@ -146,7 +146,7 @@ func TestCmd(ctx context.Context, archs []apko_types.Architecture, baseOpts ...b
 		} else if err != nil {
 			return err
 		}
-		defer bc.Close()
+		defer bc.Close(ctx)
 
 		bcs = append(bcs, bc)
 	}