Move signal handling to command context

cmd/cofidectl/main.go

@@ -4,12 +4,8 @@
 package main
 
 import (
-	"fmt"
 	"log"
 	"os"
-	"os/signal"
-	"syscall"
-	"time"
 
 	"github.com/spf13/cobra"
 
@@ -36,8 +32,7 @@ func main() {
 func run() error {
 	cmdCtx := cmdcontext.NewCommandContext(cofideConfigFile)
 	defer cmdCtx.Shutdown()
-
-	go handleSignals(cmdCtx)
+	go cmdCtx.HandleSignals()
 
 	rootCmd, err := cmd.NewRootCommand(cmdCtx).GetRootCommand()
 	if err != nil {
@@ -63,20 +58,6 @@ func run() error {
 	return rootCmd.ExecuteContext(cmdCtx.Ctx)
 }
 
-// handleSignals waits for SIGINT or SIGTERM, then triggers a clean shutdown using the command context.
-func handleSignals(cmdCtx *cmdcontext.CommandContext) {
-	shutdown := make(chan os.Signal, 1)
-	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)
-	s := <-shutdown
-	fmt.Printf("Caught %s signal, exiting\n", s.String())
-	cmdCtx.Shutdown()
-
-	// Wait for a while to allow for graceful completion of the main goroutine.
-	<-time.After(shutdownTimeoutSec * time.Second)
-	fmt.Println("Timed out waiting for shutdown")
-	os.Exit(1)
-}
-
 // getCliPlugin returns a `plugin.CliPlugin` for a CLI plugin if:
 // 1. the first CLI argument does not match a registered subcommand
 // 2. a cofidectl plugin exists with a name of cofidectl- followed by the first CLI argument

pkg/cmd/context/context.go

@@ -6,11 +6,18 @@ package context
 import (
 	"context"
 	"errors"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
 
 	"github.com/cofide/cofidectl/internal/pkg/config"
 	"github.com/cofide/cofidectl/pkg/plugin/manager"
 )
 
+const shutdownTimeoutSec = 10
+
 type CommandContext struct {
 	Ctx           context.Context
 	cancel        context.CancelCauseFunc
@@ -32,3 +39,18 @@ func (cc *CommandContext) Shutdown() {
 	}
 	cc.PluginManager.Shutdown()
 }
+
+// HandleSignals waits for SIGINT or SIGTERM, then triggers a clean shutdown using the command context.
+// It should be called from a non-main goroutine.
+func (cc *CommandContext) HandleSignals() {
+	shutdown := make(chan os.Signal, 1)
+	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)
+	s := <-shutdown
+	fmt.Printf("Caught %s signal, exiting\n", s.String())
+	cc.Shutdown()
+
+	// Wait for a while to allow for graceful completion of the main goroutine.
+	<-time.After(shutdownTimeoutSec * time.Second)
+	fmt.Println("Timed out waiting for shutdown")
+	os.Exit(1)
+}

pkg/plugin/plugin_test.go

@@ -0,0 +1,20 @@
+// Copyright 2024 Cofide Limited.
+// SPDX-License-Identifier: Apache-2.0
+
+package plugin
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsDataSourceServeCmd(t *testing.T) {
+	assert.True(t, IsDataSourceServeCmd([]string{"data-source", "serve"}))
+	assert.False(t, IsDataSourceServeCmd([]string{"trust-zone", "list"}))
+	assert.False(t, IsDataSourceServeCmd([]string{}))
+	assert.False(t, IsDataSourceServeCmd([]string{"data-source"}))
+	assert.False(t, IsDataSourceServeCmd([]string{"DATA-SOURCE", "serve"}))
+	assert.False(t, IsDataSourceServeCmd([]string{"data-source", "serve", "extra"}))
+	assert.False(t, IsDataSourceServeCmd(nil))
+}