Implement OS signal awareness for the CLI

internal/cmd/discoverworkload/discoverworkload.go

@@ -7,6 +7,8 @@ import (
 	"io"
 	"log/slog"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/spf13/cobra"
@@ -75,6 +77,8 @@ func NewCommand(ctx context.Context) *cobra.Command {
 
 			logger.Info("starting to watch for workloads", "duration", cfg.Timeout)
 
+			go gracefulShutdown(cancel)
+
 			opts := discover.NewManifestJSONProcessorFnOptions{
 				CompactOutput: cfg.CompactOutput,
 			}
@@ -143,3 +147,10 @@ func newLogger(level string, out io.Writer) (*slog.Logger, error) {
 
 	return logger, nil
 }
+
+func gracefulShutdown(cancel context.CancelFunc) {
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
+	<-quit
+	cancel()
+}

internal/discover/processor_test.go

@@ -394,6 +394,11 @@ func TestManifestJSONProcessor(t *testing.T) {
 		compact  bool
 		expected []byte
 	}{
+		"no pods": {
+			ctx:      context.TODO(),
+			input:    []corev1.Pod{},
+			expected: []byte{},
+		},
 		"initContainer only": {
 			ctx: context.TODO(),
 			input: []corev1.Pod{
@@ -461,7 +466,7 @@ func TestManifestJSONProcessor(t *testing.T) {
 			actual, err := io.ReadAll(buffer)
 
 			if testFnErr != nil {
-				t.Fatalf("processor function threw an error unexpectedly: %q", err)
+				t.Fatalf("processor function threw an error unexpectedly: %q", testFnErr)
 			}
 
 			if err != nil {
@@ -475,6 +480,60 @@ func TestManifestJSONProcessor(t *testing.T) {
 	}
 }
 
+func TestManifestJSONProcessorContextCancelled(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithCancel(context.TODO())
+	input := corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Name: "pod-1"},
+		Spec: corev1.PodSpec{
+			Containers: []corev1.Container{
+				{
+					Name:  "cname",
+					Image: "example.com/namespace/image:0.0.1",
+				},
+			},
+		},
+	}
+	expected := []byte("{\"DiscoveredImages\":[{\"Image\":\"example.com/namespace/image:0.0.1\",\"Containers\":[{\"Name\":\"cname\",\"Type\":\"Container\",\"Pod\":{\"Name\":\"pod-1\",\"Namespace\":\"\"}}]}]}\n")
+
+	testLogger := NewSlogDiscardLogger()
+	buffer := bytes.NewBuffer([]byte{})
+	opts := NewManifestJSONProcessorFnOptions{
+		CompactOutput: true,
+	}
+	fn := NewManifestJSONProcessorFn(buffer, opts)
+
+	ch := make(chan *corev1.Pod)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	var testFnErr error
+	go func() {
+		defer wg.Done()
+		testFnErr = fn(ctx, ch, testLogger)
+	}()
+
+	// If the context is cancelled prematurely, the function should
+	// still output JSON for what has already been processed.
+	ch <- &input
+	cancel()
+
+	wg.Wait()
+	actual, err := io.ReadAll(buffer)
+
+	if testFnErr != nil {
+		t.Fatalf("processor function threw an error unexpectedly: %q", testFnErr)
+	}
+
+	if err != nil {
+		t.Fatalf("unable to read output written to buffer: %q", err)
+	}
+
+	if !bytes.Equal(actual, expected) {
+		t.Fatalf("ManifestJSONPrinter processing function returned the wrong output. actual: %q expected %q", actual, expected)
+	}
+}
+
 func TestContainerProcessing(t *testing.T) {
 	t.Parallel()
 	testcases := map[string]struct {