visionA/local-tool/server/internal/firmware/shutdown_test.go

package firmware

// shutdown_test.go — M9-4.5：firmware-aware graceful shutdown helper 單元測試。
//
// 覆蓋情境：
//  1. 沒 active task → 立刻 return true、不廣播、不呼叫 RequestShutdown
//  2. 有 active task、Wait 內結束 → 廣播 1 次、return true
//  3. 有 active task、timeout → 廣播 1 次、return false
//  4. nil service → 防呆 return true
//  5. nil notifier / nil logger → 仍正常 return（不 panic）

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"testing"
	"time"
)

// fakeLifecycle 是 FirmwareLifecycle 的 test double。
type fakeLifecycle struct {
	mu sync.Mutex

	hasActive       bool
	activeTaskInfos []*ActiveTaskInfo

	// 觀察：method call 次數
	hasActiveCalls   int
	getInfoCalls     int
	requestShutdown  int
	waitForActive    int
	waitForActiveDur time.Duration

	// 控制 WaitForActiveTasks 回傳
	waitResult bool
}

func (f *fakeLifecycle) HasActiveTask() bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.hasActiveCalls++
	return f.hasActive
}

func (f *fakeLifecycle) GetActiveTaskInfo() []*ActiveTaskInfo {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.getInfoCalls++
	return f.activeTaskInfos
}

func (f *fakeLifecycle) RequestShutdown() {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.requestShutdown++
}

func (f *fakeLifecycle) WaitForActiveTasks(maxWait time.Duration) bool {
	f.mu.Lock()
	f.waitForActive++
	f.waitForActiveDur = maxWait
	res := f.waitResult
	f.mu.Unlock()
	return res
}

// fakeNotifier 觀察 BroadcastToRoom 呼叫。
type fakeNotifier struct {
	mu       sync.Mutex
	calls    int
	lastRoom string
	lastData interface{}
}

func (n *fakeNotifier) BroadcastToRoom(room string, data interface{}) {
	n.mu.Lock()
	defer n.mu.Unlock()
	n.calls++
	n.lastRoom = room
	n.lastData = data
}

// fakeLogger 觀察 Info / Warn 呼叫（內容檢查不嚴格、只看「有被呼叫」）。
type fakeLogger struct {
	mu       sync.Mutex
	infoMsgs []string
	warnMsgs []string
}

func (l *fakeLogger) Info(format string, args ...interface{}) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.infoMsgs = append(l.infoMsgs, format)
}

func (l *fakeLogger) Warn(format string, args ...interface{}) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.warnMsgs = append(l.warnMsgs, format)
}

// ──────────────────────────────────────────────────────────────────────
// 1. 沒 active task → 立刻 return true
// ──────────────────────────────────────────────────────────────────────

func TestAwaitActiveTasksOrTimeout_NoActiveTask(t *testing.T) {
	svc := &fakeLifecycle{hasActive: false}
	notifier := &fakeNotifier{}
	logger := &fakeLogger{}

	clean := AwaitActiveTasksOrTimeout(context.Background(), svc, notifier, logger)

	if !clean {
		t.Errorf("expected cleanShutdown=true when no active task, got false")
	}
	if svc.hasActiveCalls != 1 {
		t.Errorf("expected HasActiveTask called once, got %d", svc.hasActiveCalls)
	}
	if svc.requestShutdown != 0 {
		t.Errorf("expected RequestShutdown NOT called when no active task, got %d", svc.requestShutdown)
	}
	if svc.waitForActive != 0 {
		t.Errorf("expected WaitForActiveTasks NOT called when no active task, got %d", svc.waitForActive)
	}
	if notifier.calls != 0 {
		t.Errorf("expected no broadcast when no active task, got %d calls", notifier.calls)
	}
}

// ──────────────────────────────────────────────────────────────────────
// 2. 有 active task、Wait 內結束 → 廣播 1 次、return true
// ──────────────────────────────────────────────────────────────────────

func TestAwaitActiveTasksOrTimeout_ActiveTaskFinishesCleanly(t *testing.T) {
	tasks := []*ActiveTaskInfo{
		{TaskID: "upgrade-KL520-0", DeviceID: "kl520-0", Chip: ChipKL520, Direction: DirectionUpgrade, Stage: StageFlashing},
	}
	svc := &fakeLifecycle{
		hasActive:       true,
		activeTaskInfos: tasks,
		waitResult:      true, // task 在 timeout 前結束
	}
	notifier := &fakeNotifier{}
	logger := &fakeLogger{}

	clean := AwaitActiveTasksOrTimeout(context.Background(), svc, notifier, logger)

	if !clean {
		t.Errorf("expected cleanShutdown=true when WaitForActiveTasks returns true")
	}
	if svc.requestShutdown != 1 {
		t.Errorf("expected RequestShutdown called once, got %d", svc.requestShutdown)
	}
	if svc.waitForActive != 1 {
		t.Errorf("expected WaitForActiveTasks called once, got %d", svc.waitForActive)
	}
	if svc.waitForActiveDur != MaxShutdownWait {
		t.Errorf("expected wait dur=%v, got %v", MaxShutdownWait, svc.waitForActiveDur)
	}
	if notifier.calls != 1 {
		t.Errorf("expected broadcast called once, got %d", notifier.calls)
	}
	if notifier.lastRoom != ShutdownRoom {
		t.Errorf("expected broadcast room=%q, got %q", ShutdownRoom, notifier.lastRoom)
	}
	// 驗 broadcast payload schema
	payload, ok := notifier.lastData.(map[string]interface{})
	if !ok {
		t.Fatalf("broadcast payload not map[string]interface{}: %T", notifier.lastData)
	}
	if payload["type"] != ShutdownEventTypePending {
		t.Errorf("expected payload.type=%q, got %v", ShutdownEventTypePending, payload["type"])
	}
	if _, ok := payload["tasks"]; !ok {
		t.Errorf("expected payload.tasks key present")
	}
	// Minor-3：tasks 應為 minimal struct slice、不是 *ActiveTaskInfo
	broadcastTasks, ok := payload["tasks"].([]shutdownBroadcastTask)
	if !ok {
		t.Fatalf("expected payload.tasks to be []shutdownBroadcastTask (without StartTs), got %T", payload["tasks"])
	}
	if len(broadcastTasks) != 1 {
		t.Fatalf("expected 1 broadcast task, got %d", len(broadcastTasks))
	}
	if broadcastTasks[0].TaskID != "upgrade-KL520-0" {
		t.Errorf("expected broadcast task TaskID preserved, got %q", broadcastTasks[0].TaskID)
	}
}

// ──────────────────────────────────────────────────────────────────────
// 2a. Minor-3 verification：broadcast payload 不含 startTs 欄位（避免時間戳洩漏）
// ──────────────────────────────────────────────────────────────────────

func TestAwaitActiveTasksOrTimeout_BroadcastFiltersStartTs(t *testing.T) {
	// 帶非零 StartTs 的 task — 驗 broadcast 端確實過濾掉
	now := time.Now()
	tasks := []*ActiveTaskInfo{
		{
			TaskID:     "upgrade-KL720-0",
			DeviceID:   "kl720-0",
			DeviceName: "KL720 dev",
			Chip:       ChipKL720,
			Direction:  DirectionUpgrade,
			Stage:      StageFlashing,
			StartTs:    now,
			ElapsedMs:  5000,
			EtaSeconds: 30,
		},
	}
	svc := &fakeLifecycle{
		hasActive:       true,
		activeTaskInfos: tasks,
		waitResult:      true,
	}
	notifier := &fakeNotifier{}

	_ = AwaitActiveTasksOrTimeout(context.Background(), svc, notifier, nil)

	// 1. broadcast 一定要被呼叫
	if notifier.calls != 1 {
		t.Fatalf("expected broadcast called once, got %d", notifier.calls)
	}

	// 2. payload tasks 必須是 minimal struct slice
	payload, ok := notifier.lastData.(map[string]interface{})
	if !ok {
		t.Fatalf("payload not map[string]interface{}: %T", notifier.lastData)
	}
	broadcastTasks, ok := payload["tasks"].([]shutdownBroadcastTask)
	if !ok {
		t.Fatalf("expected []shutdownBroadcastTask (no StartTs), got %T — possible regression to *ActiveTaskInfo", payload["tasks"])
	}
	if len(broadcastTasks) != 1 {
		t.Fatalf("expected 1 task, got %d", len(broadcastTasks))
	}

	// 3. 序列化後驗 JSON 不含 "startTs" key（與直接 marshal *ActiveTaskInfo 對比）
	jsonBytes, err := json.Marshal(broadcastTasks[0])
	if err != nil {
		t.Fatalf("marshal broadcastTask failed: %v", err)
	}
	if strings.Contains(string(jsonBytes), "startTs") {
		t.Errorf("broadcast payload should NOT contain 'startTs' field; got JSON: %s", string(jsonBytes))
	}

	// 4. 其他欄位確實保留（taskId / deviceId / deviceName / chip / direction / stage / elapsedMs / etaSeconds 七個）
	bt := broadcastTasks[0]
	if bt.TaskID != "upgrade-KL720-0" {
		t.Errorf("TaskID lost: %q", bt.TaskID)
	}
	if bt.DeviceID != "kl720-0" {
		t.Errorf("DeviceID lost: %q", bt.DeviceID)
	}
	if bt.DeviceName != "KL720 dev" {
		t.Errorf("DeviceName lost: %q", bt.DeviceName)
	}
	if bt.Chip != ChipKL720 {
		t.Errorf("Chip lost: %q", bt.Chip)
	}
	if bt.Direction != DirectionUpgrade {
		t.Errorf("Direction lost: %q", bt.Direction)
	}
	if bt.Stage != StageFlashing {
		t.Errorf("Stage lost: %q", bt.Stage)
	}
	if bt.ElapsedMs != 5000 {
		t.Errorf("ElapsedMs lost: %d", bt.ElapsedMs)
	}
	if bt.EtaSeconds != 30 {
		t.Errorf("EtaSeconds lost: %d", bt.EtaSeconds)
	}
}

// ──────────────────────────────────────────────────────────────────────
// 2b. Minor-3 helper：toBroadcastTasks 處理 nil 與空 slice
// ──────────────────────────────────────────────────────────────────────

func TestToBroadcastTasks_NilAndEmpty(t *testing.T) {
	// 空 slice → 回空 slice（非 nil、避免 frontend null）
	out := toBroadcastTasks(nil)
	if out == nil {
		t.Errorf("toBroadcastTasks(nil) should return non-nil slice")
	}
	if len(out) != 0 {
		t.Errorf("toBroadcastTasks(nil) should return empty slice, got %d", len(out))
	}

	// 含 nil pointer 的 slice → 跳過 nil
	mixed := []*ActiveTaskInfo{
		nil,
		{TaskID: "t1", Chip: ChipKL520},
		nil,
	}
	out = toBroadcastTasks(mixed)
	if len(out) != 1 {
		t.Errorf("expected 1 valid task (nil entries filtered), got %d", len(out))
	}
	if out[0].TaskID != "t1" {
		t.Errorf("expected TaskID=t1, got %q", out[0].TaskID)
	}
}

// ──────────────────────────────────────────────────────────────────────
// 3. 有 active task、timeout → 廣播 1 次、return false、走強制 shutdown
// ──────────────────────────────────────────────────────────────────────

func TestAwaitActiveTasksOrTimeout_ActiveTaskTimeout(t *testing.T) {
	// 縮短 MaxShutdownWait 讓測試快
	origMax := MaxShutdownWait
	MaxShutdownWait = 10 * time.Millisecond
	defer func() { MaxShutdownWait = origMax }()

	svc := &fakeLifecycle{
		hasActive:       true,
		activeTaskInfos: []*ActiveTaskInfo{{TaskID: "stuck", Chip: ChipKL720}},
		waitResult:      false, // 等不到 task 結束
	}
	notifier := &fakeNotifier{}
	logger := &fakeLogger{}

	clean := AwaitActiveTasksOrTimeout(context.Background(), svc, notifier, logger)

	if clean {
		t.Errorf("expected cleanShutdown=false when WaitForActiveTasks times out")
	}
	if notifier.calls != 1 {
		t.Errorf("expected broadcast called once even on timeout, got %d", notifier.calls)
	}
	// timeout 應該觸發 Warn log（不是只有 Info）
	if len(logger.warnMsgs) == 0 {
		t.Errorf("expected at least one Warn log on timeout, got none")
	}
}

// ──────────────────────────────────────────────────────────────────────
// 4. nil service → 防呆 return true（視為「乾淨」、main 繼續走 shutdown）
// ──────────────────────────────────────────────────────────────────────

func TestAwaitActiveTasksOrTimeout_NilService(t *testing.T) {
	clean := AwaitActiveTasksOrTimeout(context.Background(), nil, nil, &fakeLogger{})
	if !clean {
		t.Errorf("expected cleanShutdown=true with nil service (defensive default)")
	}
}

// ──────────────────────────────────────────────────────────────────────
// 5. nil notifier / nil logger → 仍正常 return（不 panic）
// ──────────────────────────────────────────────────────────────────────

func TestAwaitActiveTasksOrTimeout_NilNotifierAndLogger(t *testing.T) {
	svc := &fakeLifecycle{hasActive: true, activeTaskInfos: nil, waitResult: true}

	defer func() {
		if r := recover(); r != nil {
			t.Errorf("panicked with nil notifier/logger: %v", r)
		}
	}()

	clean := AwaitActiveTasksOrTimeout(context.Background(), svc, nil, nil)
	if !clean {
		t.Errorf("expected cleanShutdown=true (waitResult=true)")
	}
	if svc.requestShutdown != 1 {
		t.Errorf("expected RequestShutdown still called with nil notifier, got %d", svc.requestShutdown)
	}
}

// ──────────────────────────────────────────────────────────────────────
// 6. 真接整合：用 real *Service 驗 helper 整條走通（不用 fake lifecycle）
// ──────────────────────────────────────────────────────────────────────

func TestAwaitActiveTasksOrTimeout_RealServiceNoActive(t *testing.T) {
	svc := NewService(&fakeLookup{drivers: map[string]UpgradeDriver{}}, FirmwareDir{Root: "/tmp"})
	notifier := &fakeNotifier{}

	clean := AwaitActiveTasksOrTimeout(context.Background(), svc, notifier, nil)
	if !clean {
		t.Errorf("expected cleanShutdown=true on fresh service with no tasks")
	}
	if notifier.calls != 0 {
		t.Errorf("expected no broadcast on fresh service, got %d", notifier.calls)
	}
}