visionA/local-tool/server/internal/api/ws/firmware_ws_test.go

package ws

// firmware_ws_test.go — M9-4 hot-fix smoke test
//
// 對稱 system_ws_integration_test.go：啟 httptest server 掛
// FirmwareProgressHandler，用 gorilla WebSocket client 真的連線進去，
// 然後 hub.BroadcastToRoom("firmware:<deviceID>", ...) 驗證 client 收到。
//
// 目的：保證 /ws/devices/:id/firmware-progress endpoint 把 client 正確
// join 到 "firmware:<deviceID>" room、且 broadcast 能送達。

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

func TestFirmwareProgressHandler_ReceivesBroadcast(t *testing.T) {
	gin.SetMode(gin.TestMode)

	hub := NewHub()
	go hub.Run()

	r := gin.New()
	r.GET("/ws/devices/:id/firmware-progress", FirmwareProgressHandler(hub))

	srv := httptest.NewServer(r)
	defer srv.Close()

	const deviceID = "kl520-0"
	room := "firmware:" + deviceID

	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/devices/" + deviceID + "/firmware-progress"
	dialer := websocket.DefaultDialer
	conn, _, err := dialer.Dial(wsURL, http.Header{})
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer conn.Close()

	_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))

	// 等 hub 吸收 Register（最多 500 ms）
	deadline := time.Now().Add(500 * time.Millisecond)
	for time.Now().Before(deadline) {
		hub.mu.RLock()
		n := len(hub.rooms[room])
		hub.mu.RUnlock()
		if n > 0 {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}

	// 廣播 firmware progress 事件
	hub.BroadcastToRoom(room, map[string]interface{}{
		"type":     "firmware:progress",
		"deviceId": deviceID,
		"phase":    "flashing",
		"percent":  42,
	})

	_, data, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("read: %v", err)
	}

	var got map[string]interface{}
	if err := json.Unmarshal(data, &got); err != nil {
		t.Fatalf("json: %v; raw=%s", err, string(data))
	}
	if got["type"] != "firmware:progress" {
		t.Errorf("wrong type: %+v", got)
	}
	if got["deviceId"] != deviceID {
		t.Errorf("wrong deviceId: %+v", got)
	}
	if got["phase"] != "flashing" {
		t.Errorf("wrong phase: %+v", got)
	}
}

// TestFirmwareProgressHandler_RoomIsolation 驗證不同 deviceID 的 client
// 不會收到對方 room 的訊息（room key 帶 deviceID 的關鍵保證）。
func TestFirmwareProgressHandler_RoomIsolation(t *testing.T) {
	gin.SetMode(gin.TestMode)

	hub := NewHub()
	go hub.Run()

	r := gin.New()
	r.GET("/ws/devices/:id/firmware-progress", FirmwareProgressHandler(hub))

	srv := httptest.NewServer(r)
	defer srv.Close()

	dialFor := func(deviceID string) *websocket.Conn {
		t.Helper()
		wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") +
			"/ws/devices/" + deviceID + "/firmware-progress"
		conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{})
		if err != nil {
			t.Fatalf("dial %s: %v", deviceID, err)
		}
		return conn
	}

	connA := dialFor("kl520-0")
	defer connA.Close()
	connB := dialFor("kl720-1")
	defer connB.Close()

	// 等兩個 room 都被 hub 吸收
	wantRooms := []string{"firmware:kl520-0", "firmware:kl720-1"}
	deadline := time.Now().Add(500 * time.Millisecond)
	for time.Now().Before(deadline) {
		ok := true
		hub.mu.RLock()
		for _, room := range wantRooms {
			if len(hub.rooms[room]) == 0 {
				ok = false
				break
			}
		}
		hub.mu.RUnlock()
		if ok {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}

	// 只 broadcast 到 kl520-0 的 room
	hub.BroadcastToRoom("firmware:kl520-0", map[string]interface{}{
		"type":     "firmware:progress",
		"deviceId": "kl520-0",
		"percent":  10,
	})

	// connA 應該收到
	_ = connA.SetReadDeadline(time.Now().Add(1 * time.Second))
	_, data, err := connA.ReadMessage()
	if err != nil {
		t.Fatalf("connA read: %v", err)
	}
	var got map[string]interface{}
	if err := json.Unmarshal(data, &got); err != nil {
		t.Fatalf("connA json: %v", err)
	}
	if got["deviceId"] != "kl520-0" {
		t.Errorf("connA got wrong deviceId: %+v", got)
	}

	// connB 不該收到 — 設短 deadline，預期 timeout
	_ = connB.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
	_, _, err = connB.ReadMessage()
	if err == nil {
		t.Errorf("connB should not have received message but did")
	}
}