chore(integration): cleanup on interrupt (#283)

johnstcn · web-flow · commit 0da135c02070 · 2024-09-06T10:48:17.000+01:00
diff --git a/integration/integration_test.go b/integration/integration_test.go
@@ -7,10 +7,12 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -203,12 +205,24 @@ func setup(ctx context.Context, t *testing.T, name, coderImg, coderVersion strin
 	require.NoError(t, err, "create test deployment")
 
 	t.Logf("created container %s\n", ctr.ID)
-	t.Cleanup(func() { // Make sure we clean up after ourselves.
-		// TODO: also have this execute if you Ctrl+C!
+	var cleanupOnce sync.Once
+	removeContainer := func() {
 		t.Logf("stopping container %s\n", ctr.ID)
 		_ = cli.ContainerRemove(ctx, ctr.ID, container.RemoveOptions{
 			Force: true,
 		})
+	}
+	// Ensure the container is cleaned up if you press Ctrl+C.
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt)
+	go func() {
+		<-sigCh
+		cleanupOnce.Do(removeContainer)
+		os.Exit(1)
+	}()
+
+	t.Cleanup(func() { // Make sure we clean up after ourselves.
+		cleanupOnce.Do(removeContainer)
 	})
 
 	err = cli.ContainerStart(ctx, ctr.ID, container.StartOptions{})
