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) } // 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)") }