package api import ( "bufio" "net" "net/http" "net/http/httptest" "strings" "testing" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "visiona-backend/internal/session" ) // newWSFixture 建立一個掛了 registerWebSocketRoutes 的 router(含 injectStaticUserContext // 模擬 AuthMiddleware 已放行)。deps 由 caller 提供以便測不同情境。 func newWSFixture(deps Deps) *gin.Engine { r := gin.New() r.Use(RequestIDMiddleware()) r.Use(injectStaticUserContext("demo-user", "")) ws := r.Group("/ws") registerWebSocketRoutes(ws, deps) return r } // TestWSInference_NoForwarder 驗證缺 Forwarder 時回 501(走 recorder,不需 hijack)。 func TestWSInference_NoForwarder(t *testing.T) { r := newWSFixture(Deps{}) w := httptest.NewRecorder() r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/ws/devices/dev1/inference", nil)) assert.Equal(t, http.StatusNotImplemented, w.Code) } // TestWSInference_TunnelDisconnected 驗證無 active session 時回 502 TUNNEL_DISCONNECTED。 func TestWSInference_TunnelDisconnected(t *testing.T) { r := newWSFixture(Deps{ SessionStore: &fakeSessionStore{}, // List 回空 Forwarder: session.NewForwarder("http://localhost:0", nil), }) w := httptest.NewRecorder() r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/ws/devices/dev1/inference", nil)) assert.Equal(t, http.StatusBadGateway, w.Code) assert.Contains(t, w.Body.String(), ErrCodeTunnelDisconnect) } // TestWSFlashProgress_NoForwarder 驗證 flash-progress route 已掛上(非 404/501-stub), // 且與 inference 共用同一個 path-agnostic handler:缺 Forwarder 時回 501(handler 內部 // 的「forwarder/session store not configured」,非舊的 registerWebSocketStubs 501)。 // // 這條與 TestWSInference_NoForwarder 對照,證明 flash-progress 走的是 // registerWebSocketRoutes(wsAuthGroup)而非 registerWebSocketStubs。 func TestWSFlashProgress_NoForwarder(t *testing.T) { r := newWSFixture(Deps{}) w := httptest.NewRecorder() r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/ws/devices/dev1/flash-progress", nil)) assert.Equal(t, http.StatusNotImplemented, w.Code) } // TestWSFlashProgress_TunnelDisconnected 驗證 flash-progress 無 active session 時回 // 502 TUNNEL_DISCONNECTED(與 inference 同一套 pickActiveSessionToken 路徑)。 func TestWSFlashProgress_TunnelDisconnected(t *testing.T) { r := newWSFixture(Deps{ SessionStore: &fakeSessionStore{}, // List 回空 Forwarder: session.NewForwarder("http://localhost:0", nil), }) w := httptest.NewRecorder() r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/ws/devices/dev1/flash-progress", nil)) assert.Equal(t, http.StatusBadGateway, w.Code) assert.Contains(t, w.Body.String(), ErrCodeTunnelDisconnect) } // TestWSFlashProgress_CoexistsWithInference 驗證 radix tree 共存: // /ws/devices/events(靜態) vs /ws/devices/:id/inference(param) vs // /ws/devices/:id/flash-progress(param)三者掛在同一 group 不 panic,且 flash-progress // 與 inference 都能各自路由到 handler(而非彼此蓋掉)。newWSFixture 建構本身若 panic // 就會 fail;這裡再各打一發確認兩條 param route 都 match 得到(回 501 = 命中 handler)。 func TestWSFlashProgress_CoexistsWithInference(t *testing.T) { r := newWSFixture(Deps{}) // 建構不 panic 即代表 radix tree 共存 OK for _, path := range []string{ "/ws/devices/dev1/inference", "/ws/devices/dev1/flash-progress", } { w := httptest.NewRecorder() r.ServeHTTP(w, httptest.NewRequest(http.MethodGet, path, nil)) // 命中 handler(缺 Forwarder → 501)而非 404(route 沒掛上)。 assert.Equal(t, http.StatusNotImplemented, w.Code, "path %s 應命中 WS proxy handler", path) } } // TestWSInference_HappyPath_BidirectionalPipe 是端到端的 WS forward 測試: // // browser(test client) ──► api-server handler(hijack) ──► Forwarder.OpenStream // ──► fake remote-proxy(raw, hijack + 200 Connected) ──► fake local agent(101 + echo) // // 驗證: // - handler 回 101 給 browser(升級成功) // - 升級後 browser → agent 的 bytes 能雙向 pipe(送 ping 收 echo) // // 用真 httptest.NewServer(recorder 不支援 Hijack)。 func TestWSInference_HappyPath_BidirectionalPipe(t *testing.T) { const clientMsg = "ping-from-browser" // 1. 假 remote-proxy:hijack → 200 Connected → 模擬 local agent(讀 upgrade → 回 101 → echo) fakeProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hj, ok := w.(http.Hijacker) require.True(t, ok) conn, _, err := hj.Hijack() require.NoError(t, err) defer conn.Close() _, _ = conn.Write([]byte("HTTP/1.1 200 Connected\r\n\r\n")) // 模擬 local agent:讀 upgrade request → 回 101 → 之後把收到的 bytes echo 回去 br := bufio.NewReader(conn) req, err := http.ReadRequest(br) if err != nil { return } assert.True(t, strings.EqualFold(req.Header.Get("Upgrade"), "websocket")) _, _ = conn.Write([]byte( "HTTP/1.1 101 Switching Protocols\r\n" + "Upgrade: websocket\r\nConnection: Upgrade\r\n" + "Sec-WebSocket-Accept: x\r\n\r\n", )) // echo:讀 browser 送來的 bytes(透過 handler pipe)→ 原樣寫回 buf := make([]byte, len(clientMsg)) n, _ := br.Read(buf) if n > 0 { _, _ = conn.Write(buf[:n]) } })) defer fakeProxy.Close() // 2. api-server router(有 active session token → pickActiveSessionToken 命中) r := newWSFixture(Deps{ SessionStore: &fakeSessionStore{sessions: []*session.Summary{ {Token: "vAc_live", UserID: "demo-user"}, }}, Forwarder: session.NewForwarder(fakeProxy.URL, nil), }) apiServer := httptest.NewServer(r) defer apiServer.Close() // 3. browser 端:raw TCP dial + 送 WS upgrade request u := strings.TrimPrefix(apiServer.URL, "http://") cli, err := net.Dial("tcp", u) require.NoError(t, err) defer cli.Close() upgrade := "GET /ws/devices/dev1/inference HTTP/1.1\r\n" + "Host: " + u + "\r\n" + "Upgrade: websocket\r\n" + "Connection: Upgrade\r\n" + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" + "Sec-WebSocket-Version: 13\r\n\r\n" _, err = cli.Write([]byte(upgrade)) require.NoError(t, err) // 讀 101 response cbr := bufio.NewReader(cli) resp, err := http.ReadResponse(cbr, nil) require.NoError(t, err) assert.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode, "handler 應回 101 給 browser") // 升級後:送一段 bytes,期望 echo 回來(驗證雙向 pipe) _, err = cli.Write([]byte(clientMsg)) require.NoError(t, err) echo := make([]byte, len(clientMsg)) _, err = cbr.Read(echo) require.NoError(t, err) assert.Equal(t, clientMsg, string(echo), "升級後 bytes 應雙向 pipe(收到 echo)") }