fix(agent): Make Ctrl+C actually work: Do not use context.TODO

agent/cloudinit/cloudinit.go

@@ -1,4 +1,5 @@
 // Copyright 2021 VMware, Inc. All Rights Reserved.
+// Copyright 2026 Platform9, Inc. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
 package cloudinit
@@ -43,7 +44,7 @@ type Files struct {
 //   - parse the script to get the cloudinit data
 //   - execute the write_files directive
 //   - execute the run_cmd directive
-func (se ScriptExecutor) Execute(bootstrapScript string) error {
+func (se ScriptExecutor) Execute(ctx context.Context, bootstrapScript string) error {
 	cloudInitData := bootstrapConfig{}
 	if err := yaml.Unmarshal([]byte(bootstrapScript), &cloudInitData); err != nil {
 		return errors.Wrapf(err, "error parsing write_files action: %s", bootstrapScript)
@@ -74,7 +75,7 @@ func (se ScriptExecutor) Execute(bootstrapScript string) error {
 	}
 
 	for _, cmd := range cloudInitData.CommandsToExecute {
-		err := se.RunCmdExecutor.RunCmd(context.TODO(), cmd)
+		err := se.RunCmdExecutor.RunCmd(ctx, cmd)
 		if err != nil {
 			return errors.Wrap(err, fmt.Sprintf("Error running the command %s", cmd))
 		}

agent/cloudinit/cloudinit_integration_test.go

@@ -1,4 +1,5 @@
 // Copyright 2021 VMware, Inc. All Rights Reserved.
+// Copyright 2026 Platform9, Inc. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
 package cloudinit_test
@@ -40,7 +41,7 @@ var _ = Describe("CloudinitIntegration", func() {
 		}
 	})
 
-	It("should be able to write files and execute commands", func() {
+	It("should be able to write files and execute commands", func(ctx SpecContext) {
 		fileName := path.Join(workDir, "file-1.txt")
 		fileOriginContent := "some-content-1"
 		fileNewContent := " run cmd"
@@ -51,24 +52,24 @@ content: %s
 runCmd:
 - echo -n '%s' > %s`, fileName, fileOriginContent, fileNewContent, fileName)
 
-		err := scriptExecutor.Execute(cloudInitScript)
+		err := scriptExecutor.Execute(ctx, cloudInitScript)
 		Expect(err).ToNot(HaveOccurred())
 
 		fileContents, errFileContents := os.ReadFile(fileName)
 		Expect(errFileContents).ToNot(HaveOccurred())
 		Expect(string(fileContents)).To(Equal(fileNewContent))
 	})
 
-	It("should return error if execute commands fails", func() {
+	It("should return error if execute commands fails", func(ctx SpecContext) {
 		cloudInitScript := `
 runCmd:
 - foo`
 
-		err := scriptExecutor.Execute(cloudInitScript)
+		err := scriptExecutor.Execute(ctx, cloudInitScript)
 		Expect(err).To(HaveOccurred())
 	})
 
-	It("should be able to write files with the correct permissions and in append mode", func() {
+	It("should be able to write files with the correct permissions and in append mode", func(ctx SpecContext) {
 		fileName := path.Join(workDir, "file-2.txt")
 		fileOriginContent := "some-content-2"
 		fileAppendContent := "some-content-append-2"
@@ -84,7 +85,7 @@ runCmd:
   content: %s
   append: %v`, fileName, strconv.FormatInt(int64(filePermission), 8), fileAppendContent, isAppend)
 
-		err = scriptExecutor.Execute(cloudInitScript)
+		err = scriptExecutor.Execute(ctx, cloudInitScript)
 		Expect(err).ToNot(HaveOccurred())
 
 		fileContents, errFileContents := os.ReadFile(fileName)
@@ -96,7 +97,7 @@ runCmd:
 		Expect(stats.Mode()).To(Equal(fs.FileMode(filePermission)))
 	})
 
-	It("should be able to write encoded content", func() {
+	It("should be able to write encoded content", func(ctx SpecContext) {
 		fileName := path.Join(workDir, "file-3.txt")
 		fileContent := "some-content-3"
 		fileBase64Content := base64.StdEncoding.EncodeToString([]byte(fileContent))
@@ -106,15 +107,15 @@ runCmd:
   content: %s
   encoding: base64`, fileName, fileBase64Content)
 
-		err := scriptExecutor.Execute(cloudInitScript)
+		err := scriptExecutor.Execute(ctx, cloudInitScript)
 		Expect(err).ToNot(HaveOccurred())
 
 		fileContents, err := os.ReadFile(fileName)
 		Expect(err).ToNot(HaveOccurred())
 		Expect(string(fileContents)).To(Equal(fileContent))
 	})
 
-	It("should be able to write gziped content", func() {
+	It("should be able to write gziped content", func(ctx SpecContext) {
 		fileName := path.Join(workDir, "file-4.txt")
 		fileContent := "some-content-4"
 		fileGzipContent, err := common.GzipData([]byte(fileContent))
@@ -126,15 +127,15 @@ runCmd:
   encoding: gzip+base64
   content: %s`, fileName, fileGzipBase64Content)
 
-		err = scriptExecutor.Execute(cloudInitScript)
+		err = scriptExecutor.Execute(ctx, cloudInitScript)
 		Expect(err).ToNot(HaveOccurred())
 
 		fileContents, err := os.ReadFile(fileName)
 		Expect(err).ToNot(HaveOccurred())
 		Expect(string(fileContents)).To(Equal(fileContent))
 	})
 
-	It("should be able to write template content", func() {
+	It("should be able to write template content", func(ctx SpecContext) {
 		fileName := path.Join(workDir, "file-5.txt")
 		fileContent := "The default interface name is {{ .DefaultNetworkInterfaceName }} "
 		replacedFileContent := "The default interface name is eth0"
@@ -143,18 +144,18 @@ runCmd:
 - path: %s
   content: %s`, fileName, fileContent)
 
-		err := scriptExecutor.Execute(cloudInitScript)
+		err := scriptExecutor.Execute(ctx, cloudInitScript)
 		Expect(err).ToNot(HaveOccurred())
 
 		fileContents, err := os.ReadFile(fileName)
 		Expect(err).ToNot(HaveOccurred())
 		Expect(string(fileContents)).To(Equal(replacedFileContent))
 	})
 
-	It("should return error for invalid template content", func() {
+	It("should return error for invalid template content", func(ctx SpecContext) {
 		cloudInitScript := "invalid-content"
 
-		err := scriptExecutor.Execute(cloudInitScript)
+		err := scriptExecutor.Execute(ctx, cloudInitScript)
 		Expect(err).To(HaveOccurred())
 	})
 

agent/cloudinit/cloudinit_test.go

@@ -1,4 +1,5 @@
 // Copyright 2021 VMware, Inc. All Rights Reserved.
+// Copyright 2026 Platform9, Inc. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
 package cloudinit_test
@@ -56,7 +57,7 @@ runCmd:
 			Expect(err).NotTo(HaveOccurred())
 		})
 
-		It("should write files successfully", func() {
+		It("should write files successfully", func(ctx SpecContext) {
 			fileDir1 := path.Join(workDir, "dir1")
 			fileName1 := path.Join(fileDir1, "file1.txt")
 			fileContent1 := "some-unique-content-1"
@@ -77,7 +78,7 @@ runCmd:
   append: true
   encoding: %s`, fileName1, fileContent1, fileName2, fileBase64Content, permissions, encoding)
 
-			err = scriptExecutor.Execute(bootstrapSecretUnencoded)
+			err = scriptExecutor.Execute(ctx, bootstrapSecretUnencoded)
 			Expect(err).NotTo(HaveOccurred())
 
 			Expect(fakeFileWriter.MkdirIfNotExistsCallCount()).To(Equal(2))
@@ -102,55 +103,55 @@ runCmd:
 
 		})
 
-		It("should error out when an invalid yaml is passed", func() {
-			err := scriptExecutor.Execute("invalid yaml")
+		It("should error out when an invalid yaml is passed", func(ctx SpecContext) {
+			err := scriptExecutor.Execute(ctx, "invalid yaml")
 
 			Expect(err).To(HaveOccurred())
 			Expect(err.Error()).To(ContainSubstring("error parsing write_files action"))
 		})
 
-		It("should error out when there is not enough permission to mkdir", func() {
+		It("should error out when there is not enough permission to mkdir", func(ctx SpecContext) {
 			fakeFileWriter.MkdirIfNotExistsReturns(errors.New("not enough permissions"))
 
-			err := scriptExecutor.Execute(defaultBootstrapSecret)
+			err := scriptExecutor.Execute(ctx, defaultBootstrapSecret)
 
 			Expect(err).To(HaveOccurred())
 			Expect(err.Error()).To(ContainSubstring("not enough permissions"))
 			Expect(fakeFileWriter.WriteToFileCallCount()).To(Equal(0))
 
 		})
 
-		It("should error out write to file failes", func() {
+		It("should error out write to file failes", func(ctx SpecContext) {
 			fakeFileWriter.WriteToFileReturns(errors.New("cannot write to file"))
 
-			err := scriptExecutor.Execute(defaultBootstrapSecret)
+			err := scriptExecutor.Execute(ctx, defaultBootstrapSecret)
 
 			Expect(err).To(HaveOccurred())
 			Expect(err.Error()).To(ContainSubstring("cannot write to file"))
 		})
 
-		It("run the command given in the runCmd directive", func() {
-			err := scriptExecutor.Execute(defaultBootstrapSecret)
+		It("run the command given in the runCmd directive", func(ctx SpecContext) {
+			err := scriptExecutor.Execute(ctx, defaultBootstrapSecret)
 			Expect(err).NotTo(HaveOccurred())
 
 			Expect(fakeCmdExecutor.RunCmdCallCount()).To(Equal(1))
 			_, cmd := fakeCmdExecutor.RunCmdArgsForCall(0)
 			Expect(cmd).To(Equal("echo 'some run command'"))
 		})
 
-		It("should not invoke the runCmd or writeFiles directive when absent", func() {
+		It("should not invoke the runCmd or writeFiles directive when absent", func(ctx SpecContext) {
 
-			err := scriptExecutor.Execute("")
+			err := scriptExecutor.Execute(ctx, "")
 			Expect(err).NotTo(HaveOccurred())
 
 			Expect(fakeCmdExecutor.RunCmdCallCount()).To(Equal(0))
 			Expect(fakeFileWriter.MkdirIfNotExistsCallCount()).To(Equal(0))
 			Expect(fakeFileWriter.WriteToFileCallCount()).To(Equal(0))
 		})
 
-		It("should error out when command execution fails", func() {
+		It("should error out when command execution fails", func(ctx SpecContext) {
 			fakeCmdExecutor.RunCmdReturns(errors.New("command execution failed"))
-			err := scriptExecutor.Execute(defaultBootstrapSecret)
+			err := scriptExecutor.Execute(ctx, defaultBootstrapSecret)
 			Expect(err).To(HaveOccurred())
 
 			Expect(fakeCmdExecutor.RunCmdCallCount()).To(Equal(1))

agent/host_agent_suite_test.go

@@ -20,13 +20,9 @@ import (
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 	"github.com/onsi/gomega/gexec"
-	infrastructurev1beta1 "github.com/vmware-tanzu/cluster-api-provider-bringyourownhost/apis/infrastructure/v1beta1"
 	"github.com/vmware-tanzu/cluster-api-provider-bringyourownhost/test/e2e"
-	certv1 "k8s.io/api/certificates/v1"
 	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/runtime"
 	clientset "k8s.io/client-go/kubernetes"
-	clusterv1 "sigs.k8s.io/cluster-api/api/v1beta1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/envtest"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
@@ -56,6 +52,11 @@ var (
 	// hostNameSuffix is a per-run Unix timestamp appended to os.Hostname() to make
 	// container names unique across test runs, preventing conflicts from stale containers.
 	hostNameSuffix string
+	// suiteCtx spans the whole suite. A ByoHostRunner's context is created in a
+	// BeforeEach but its docker calls continue through the spec and into AfterEach,
+	// so it cannot be rooted on a per-node SpecContext, which Ginkgo cancels as
+	// soon as the node body returns.
+	suiteCtx context.Context
 )
 
 const (
@@ -64,7 +65,13 @@ const (
 )
 
 func TestHostAgent(t *testing.T) {
+	suiteCtx = t.Context()
 	RegisterFailHandler(Fail)
+	// Gomega drops its default Eventually timeout once an assertion carries a
+	// context, leaving a failing wait to block until the whole suite times out.
+	// The specs here pass a context to propagate it, not to bound the wait, so
+	// keep the default timeout in force.
+	EnforceDefaultTimeoutsWhenUsingContexts()
 	RunSpecs(t, "Agent Suite")
 }
 
@@ -93,21 +100,7 @@ var _ = BeforeSuite(func() {
 	Expect(err).NotTo(HaveOccurred())
 	Expect(cfg).NotTo(BeNil())
 
-	scheme := runtime.NewScheme()
-
-	err = infrastructurev1beta1.AddToScheme(scheme)
-	Expect(err).NotTo(HaveOccurred())
-
-	err = corev1.AddToScheme(scheme)
-	Expect(err).NotTo(HaveOccurred())
-
-	err = clusterv1.AddToScheme(scheme)
-	Expect(err).NotTo(HaveOccurred())
-
-	err = certv1.AddToScheme(scheme)
-	Expect(err).NotTo(HaveOccurred())
-
-	k8sClient, err = client.New(cfg, client.Options{Scheme: scheme})
+	k8sClient, err = client.New(cfg, client.Options{Scheme: newScheme()})
 	Expect(err).NotTo(HaveOccurred())
 
 	clientSet = clientset.NewForConfigOrDie(cfg)