From 3cd7f84261f2c502a710851fad8f3f3b70f433f3 Mon Sep 17 00:00:00 2001 From: Mark Goddard Date: Wed, 4 Dec 2024 09:59:05 +0000 Subject: [PATCH 1/2] Add unit test for IsDataSourceServeCmd Co-Authored-By: Nial Daly --- pkg/plugin/plugin_test.go | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 pkg/plugin/plugin_test.go diff --git a/pkg/plugin/plugin_test.go b/pkg/plugin/plugin_test.go new file mode 100644 index 0000000..9ec1a4f --- /dev/null +++ b/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)) +} From 406103951c6784fcdc66462974398cdf361542ab Mon Sep 17 00:00:00 2001 From: Mark Goddard Date: Wed, 4 Dec 2024 15:04:18 +0000 Subject: [PATCH 2/2] Move signal handling to command context This allows it to be used by plugins. --- cmd/cofidectl/main.go | 21 +-------------------- pkg/cmd/context/context.go | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 20 deletions(-) diff --git a/cmd/cofidectl/main.go b/cmd/cofidectl/main.go index 0597cbd..d5cc8db 100644 --- a/cmd/cofidectl/main.go +++ b/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 diff --git a/pkg/cmd/context/context.go b/pkg/cmd/context/context.go index 2626f8c..cf5fe15 100644 --- a/pkg/cmd/context/context.go +++ b/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) +}