ja4-platform/shared/go/ja4common/shutdown/handler_test.go

package shutdown

import (
	"context"
	"errors"
	"sync/atomic"
	"syscall"
	"testing"
	"time"
)

type mockLogger struct {
	infoMsgs  []string
	errorMsgs []string
}

func (m *mockLogger) Info(msg string)          { m.infoMsgs = append(m.infoMsgs, msg) }
func (m *mockLogger) Error(msg string, _ error) { m.errorMsgs = append(m.errorMsgs, msg) }

func TestHandle_RunsHooks(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	logger := &mockLogger{}
	var hookCalled int32

	hooks := []Hook{
		{
			Name: "test-hook",
			Fn: func() error {
				atomic.StoreInt32(&hookCalled, 1)
				return nil
			},
		},
	}

	done := make(chan struct{})
	go func() {
		Handle(ctx, cancel, hooks, logger)
		close(done)
	}()

	// Send SIGTERM to trigger shutdown
	time.Sleep(50 * time.Millisecond)
	p, _ := syscall.Getpid(), 0
	syscall.Kill(p, syscall.SIGTERM)

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("Handle did not return within timeout")
	}

	if atomic.LoadInt32(&hookCalled) != 1 {
		t.Error("hook was not called")
	}
}

func TestHandle_HookError_ContinuesOtherHooks(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	logger := &mockLogger{}
	var secondCalled int32

	hooks := []Hook{
		{
			Name: "failing-hook",
			Fn:   func() error { return errors.New("hook error") },
		},
		{
			Name: "second-hook",
			Fn: func() error {
				atomic.StoreInt32(&secondCalled, 1)
				return nil
			},
		},
	}

	done := make(chan struct{})
	go func() {
		Handle(ctx, cancel, hooks, logger)
		close(done)
	}()

	time.Sleep(50 * time.Millisecond)
	syscall.Kill(syscall.Getpid(), syscall.SIGTERM)

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("Handle did not return within timeout")
	}

	if atomic.LoadInt32(&secondCalled) != 1 {
		t.Error("second hook should still run after first hook error")
	}
	if len(logger.errorMsgs) == 0 {
		t.Error("error should be logged for failing hook")
	}
}

func TestHandle_ContextCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	logger := &mockLogger{}
	var hookCalled int32

	hooks := []Hook{
		{
			Name: "ctx-hook",
			Fn: func() error {
				atomic.StoreInt32(&hookCalled, 1)
				return nil
			},
		},
	}

	done := make(chan struct{})
	go func() {
		Handle(ctx, cancel, hooks, logger)
		close(done)
	}()

	// Cancel context directly instead of sending signal
	time.Sleep(50 * time.Millisecond)
	cancel()

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("Handle did not return within timeout after context cancel")
	}

	if atomic.LoadInt32(&hookCalled) != 1 {
		t.Error("hook should run on context cancel")
	}
}