From 02a5d7630d84ed043b0f55a4f0b66efb237e59b5 Mon Sep 17 00:00:00 2001
From: Brian Drennan <brian.drennan@cloudzero.com>
Date: Fri, 6 Sep 2024 09:20:14 -0700
Subject: [PATCH] add context object with cancel

---
 cmd/cloudzero-agent-validator/main.go | 10 +++++++---
 pkg/cmd/config/command.go             |  5 +++--
 pkg/k8s/services.go                   |  4 ++--
 pkg/k8s/services_test.go              |  5 ++++-
 4 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/cmd/cloudzero-agent-validator/main.go b/cmd/cloudzero-agent-validator/main.go
index 9ede27a..acf315e 100644
--- a/cmd/cloudzero-agent-validator/main.go
+++ b/cmd/cloudzero-agent-validator/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/signal"
@@ -17,7 +18,7 @@ import (
 )
 
 func main() {
-	ctrlCHandler()
+	ctx := ctrlCHandler()
 
 	app := &cli.App{
 		Name:     build.AppName,
@@ -36,7 +37,7 @@ func main() {
 	}
 
 	app.Commands = append(app.Commands,
-		configcmd.NewCommand(),
+		configcmd.NewCommand(ctx),
 		diagcmd.NewCommand(),
 	)
 
@@ -46,11 +47,14 @@ func main() {
 	}
 }
 
-func ctrlCHandler() {
+func ctrlCHandler() context.Context {
+	ctx, cancel := context.WithCancel(context.Background())
 	stopCh := make(chan os.Signal, 1)
 	signal.Notify(stopCh, os.Interrupt)
 	go func() {
 		<-stopCh
+		cancel()
 		os.Exit(0)
 	}()
+	return ctx
 }
diff --git a/pkg/cmd/config/command.go b/pkg/cmd/config/command.go
index df2a357..36239dc 100644
--- a/pkg/cmd/config/command.go
+++ b/pkg/cmd/config/command.go
@@ -1,6 +1,7 @@
 package config
 
 import (
+	"context"
 	_ "embed"
 	"html/template"
 	"os"
@@ -23,7 +24,7 @@ var (
 	configAlias = []string{"f"}
 )
 
-func NewCommand() *cli.Command {
+func NewCommand(ctx context.Context) *cli.Command {
 	cmd := &cli.Command{
 		Name:  "config",
 		Usage: "configuration utility commands",
@@ -94,7 +95,7 @@ func NewCommand() *cli.Command {
 						return errors.Wrap(err, "creating clientset")
 					}
 
-					return k8s.ListServices(clientset)
+					return k8s.ListServices(ctx, clientset)
 				},
 			},
 		},
diff --git a/pkg/k8s/services.go b/pkg/k8s/services.go
index 6704a4f..03b86fb 100644
--- a/pkg/k8s/services.go
+++ b/pkg/k8s/services.go
@@ -10,9 +10,9 @@ import (
 )
 
 // ListServices lists all Kubernetes services in all namespaces
-func ListServices(clientset kubernetes.Interface) error {
+func ListServices(ctx context.Context, clientset kubernetes.Interface) error {
     // List all services in all namespaces
-    services, err := clientset.CoreV1().Services("").List(context.TODO(), metav1.ListOptions{})
+		services, err := clientset.CoreV1().Services("").List(ctx, metav1.ListOptions{})
     if err != nil {
         return errors.Wrap(err, "listing services")
     }
diff --git a/pkg/k8s/services_test.go b/pkg/k8s/services_test.go
index bcc11f4..7f3bfeb 100644
--- a/pkg/k8s/services_test.go
+++ b/pkg/k8s/services_test.go
@@ -1,6 +1,7 @@
 package k8s_test
 
 import (
+		"context"
     "testing"
 
     "github.com/stretchr/testify/assert"
@@ -29,7 +30,9 @@ func TestListServices(t *testing.T) {
 		},
 	)
 
+	ctx, _ := context.WithCancel(context.Background())
+
 	// Test listing services
-	err := k8s.ListServices(clientset)
+	err := k8s.ListServices(ctx, clientset)
 	assert.NoError(t, err)
 }