From fd2828f8bc99433d623a89933bf387ded0de8986 Mon Sep 17 00:00:00 2001
From: Cian Johnston <cian@coder.com>
Date: Fri, 6 Sep 2024 10:29:04 +0100
Subject: [PATCH] chore(integration): cleanup on interrupt

---
 integration/integration_test.go | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/integration/integration_test.go b/integration/integration_test.go
index b7132d18..c98a45eb 100644
--- 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{})