Fixed tests for GracefulShutdown

internal/helpers/graceful_shutdown.go

@@ -2,29 +2,40 @@ package helpers
 
 import (
 	"context"
-	"fmt"
 	"os"
 	"os/signal"
 	"syscall"
 )
 
 var notifyFn = signal.Notify
+var sigintFix = func() {
+	// fix prints after "^C"
+	println("") // nolint:forbidigo
+}
 
 func GracefulShutdown(ctx context.Context, shutdownFunc func(ctx context.Context) error) error {
+	if done := waiteSignal(ctx); done {
+		return nil
+	}
+
+	return shutdownFunc(ctx)
+}
+
+func waiteSignal(ctx context.Context) bool {
 	stop := make(chan os.Signal, 1)
+
 	notifyFn(stop, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
 
+	defer close(stop)
+
 	select {
 	case sig := <-stop:
 		if sig == syscall.SIGINT {
-			// fix prints after "^C"
-			fmt.Println("") // nolint:forbidigo
+			sigintFix()
 		}
 	case <-ctx.Done():
-		return nil
+		return true
 	}
 
-	close(stop)
-
-	return shutdownFunc(ctx)
+	return false
 }

internal/helpers/graceful_shutdown_internal_test.go

@@ -7,6 +7,7 @@ import (
 	"sync"
 	"syscall"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -22,10 +23,6 @@ func (e *Env) Go(action func()) {
 		defer e.wg.Done()
 		action()
 	}()
-
-	for _, f := range e.afterAll {
-		f()
-	}
 }
 
 func (e *Env) CheckAfterAll(action func()) {
@@ -37,6 +34,9 @@ func WithGoroutines(test func(t *testing.T, env Env)) func(t *testing.T) {
 		env := Env{wg: &sync.WaitGroup{}}
 		test(t, env)
 		env.wg.Wait()
+		for _, f := range env.afterAll {
+			f()
+		}
 	}
 }
 
@@ -61,32 +61,82 @@ func TestGracefulShutdown(t *testing.T) {
 		cancel()
 	}))
 
-	t.Run("shutdown after system signal", WithGoroutines(func(t *testing.T, env Env) {
-		t.Skip()
-		var systemSig chan<- os.Signal
+	t.Run("shutdown after system signal", func(t *testing.T) {
+		tests := []struct {
+			name   string
+			signal os.Signal
+		}{
+			{
+				name:   "SIGINT",
+				signal: syscall.SIGINT,
+			},
+			{
+				name:   "SIGTERM",
+				signal: syscall.SIGTERM,
+			},
+			{
+				name:   "SIGHUP",
+				signal: syscall.SIGHUP,
+			},
+		}
+
+		for _, testCase := range tests {
+			t.Run(testCase.name, WithGoroutines(func(t *testing.T, env Env) {
+				var systemSig chan<- os.Signal
+				notifyFn = func(c chan<- os.Signal, sig ...os.Signal) {
+					systemSig = c
+				}
+
+				t.Cleanup(func() {
+					notifyFn = signal.Notify
+				})
+
+				called := false
+				env.Go(func() {
+					err := GracefulShutdown(context.Background(), func(ctx context.Context) error {
+						called = true
+
+						return nil
+					})
+					assert.NoError(t, err)
+				})
+
+				<-time.After(50 * time.Millisecond)
+				systemSig <- testCase.signal
+
+				env.CheckAfterAll(func() {
+					assert.True(t, called)
+				})
+			}))
+		}
+	})
 
+	t.Run("apply additional ui fix for SIGINT signal", WithGoroutines(func(t *testing.T, env Env) {
+		var systemSig chan<- os.Signal
 		notifyFn = func(c chan<- os.Signal, sig ...os.Signal) {
 			systemSig = c
 		}
+		called := false
+		sigintFix = func() {
+			called = true
+		}
 
 		t.Cleanup(func() {
 			notifyFn = signal.Notify
 		})
 
-		called := false
 		env.Go(func() {
 			err := GracefulShutdown(context.Background(), func(ctx context.Context) error {
-				called = true
-
 				return nil
 			})
 			assert.NoError(t, err)
 		})
 
+		<-time.After(50 * time.Millisecond)
+		systemSig <- syscall.SIGINT
+
 		env.CheckAfterAll(func() {
 			assert.True(t, called)
 		})
-
-		systemSig <- syscall.SIGTERM
 	}))
 }