visionA/visionA-backend/internal/api/camera_ws_test.go

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