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:", ...) 驗證 client 收到。 // // 目的:保證 /ws/devices/:id/firmware-progress endpoint 把 client 正確 // join 到 "firmware:" 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") } }