package ws_test import ( "encoding/json" "errors" "net/http" "net/http/httptest" "strings" "sync" "sync/atomic" "testing" "time" "github.com/EthanCodeCraft/xlgo-core/ws" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" ) func init() { gin.SetMode(gin.TestMode) } // newWSTestServer 启动一个带 ws.Handle 的 httptest server,返回 server 与 dialer。 // handler 可为 nil(用默认空处理器)。 func newWSTestServer(t *testing.T, handler ws.Handler) (*httptest.Server, *websocket.Dialer) { t.Helper() r := gin.New() if handler == nil { handler = &ws.DefaultHandler{} } r.GET("/ws", ws.Handle(handler)) srv := httptest.NewServer(r) t.Cleanup(srv.Close) dialer := &websocket.Dialer{HandshakeTimeout: 2 * time.Second} return srv, dialer } func wsURL(srv *httptest.Server) string { return "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" } // dial 连接并返回客户端 conn。 func dial(t *testing.T, srv *httptest.Server, dialer *websocket.Dialer) *websocket.Conn { t.Helper() c, _, err := dialer.Dial(wsURL(srv), nil) if err != nil { t.Fatalf("dial: %v", err) } t.Cleanup(func() { c.Close() }) return c } // ===== C2b:Close 与并发 Send 不 panic ===== // 回归 C2b:并发 Close 与 Send 不能 send-on-closed panic。 // 旧实现 Close 同时 close(c.send),并发 Send 的 select 伪随机选中 send 分支即 panic。 func TestConnectionCloseConcurrentSendNoPanic(t *testing.T) { srv, dialer := newWSTestServer(t, nil) client := dial(t, srv, dialer) // 服务端 conn 经 Handle 持有;通过 Hub 注册拿到服务端 Connection。 hub := ws.NewHub() go hub.Run() t.Cleanup(func() { /* Hub.Run 无退出,靠进程结束 */ }) // 用 HandleFunc 包装拿到服务端 Connection。 var srvConn *ws.Connection var got atomic.Value // *ws.Connection r := gin.New() r.GET("/ws", ws.HandleFunc(func(conn *ws.Connection, message []byte) { got.Store(conn) })) srv2 := httptest.NewServer(r) t.Cleanup(srv2.Close) c2, _, err := dialer.Dial("ws"+strings.TrimPrefix(srv2.URL, "http")+"/ws", nil) if err != nil { t.Fatalf("dial: %v", err) } defer c2.Close() _ = client // 发一条消息触发服务端 OnMessage,拿到 srvConn。 if err := c2.WriteMessage(websocket.TextMessage, []byte("hi")); err != nil { t.Fatalf("write: %v", err) } // 等待服务端拿到连接。 deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { if v := got.Load(); v != nil { srvConn = v.(*ws.Connection) break } time.Sleep(5 * time.Millisecond) } if srvConn == nil { t.Fatal("timeout waiting for server-side Connection") } // 并发:一个 goroutine 反复 Send,主 goroutine Close。 var panicked atomic.Value var wg sync.WaitGroup wg.Add(1) go func() { defer wg.Done() defer func() { if r := recover(); r != nil { panicked.Store(r) } }() for i := 0; i < 2000; i++ { // Send 在 Close 后会返回 error,正常;不应 panic。 _ = srvConn.Send([]byte("x")) } }() // 给 Send 一点启动时间再 Close,制造最大竞态窗口。 time.Sleep(2 * time.Millisecond) srvConn.Close() wg.Wait() if p := panicked.Load(); p != nil { t.Fatalf("concurrent Close/Send panicked (C2b send-on-closed): %v", p) } } // 回归 C2b:Close 后 Send 返回 error 而非 panic。 func TestSendAfterCloseReturnsError(t *testing.T) { srv, dialer := newWSTestServer(t, nil) c := dial(t, srv, dialer) defer c.Close() // 通过 HandleFunc 拿服务端 conn。 var got atomic.Value r := gin.New() r.GET("/ws", ws.HandleFunc(func(conn *ws.Connection, message []byte) { got.Store(conn) })) srv2 := httptest.NewServer(r) t.Cleanup(srv2.Close) c2, _, err := dialer.Dial("ws"+strings.TrimPrefix(srv2.URL, "http")+"/ws", nil) if err != nil { t.Fatalf("dial: %v", err) } defer c2.Close() if err := c2.WriteMessage(websocket.TextMessage, []byte("hi")); err != nil { t.Fatalf("write: %v", err) } deadline := time.Now().Add(2 * time.Second) var srvConn *ws.Connection for time.Now().Before(deadline) { if v := got.Load(); v != nil { srvConn = v.(*ws.Connection) break } time.Sleep(5 * time.Millisecond) } if srvConn == nil { t.Fatal("timeout") } srvConn.Close() // Close 后 Send 必须返回 error(非 panic)。 if err := srvConn.Send([]byte("after-close")); err == nil { t.Error("Send after Close should return error, got nil") } } // ===== C2a:Hub 广播不死锁 ===== // 回归 C2a:向含"已关闭"连接的 Hub 广播不能死锁,且失败连接被行内移除。 // 旧实现:broadcast 中 conn.Send 失败 → h.unregister <- conn(向自己发,无消费者)→ 永久阻塞。 // 这里直接 Close 服务端连接(closeChan 关闭),使 Send 走 closeChan 分支返回 error, // 触发 Hub 行内 delete + Close。 func TestHubBroadcastDeadConnectionNoDeadlock(t *testing.T) { hub := ws.NewHub() go hub.Run() var got atomic.Value r := gin.New() r.GET("/ws", ws.HandleFunc(func(conn *ws.Connection, message []byte) { got.Store(conn) })) srv2 := httptest.NewServer(r) t.Cleanup(srv2.Close) dialer := &websocket.Dialer{HandshakeTimeout: 2 * time.Second} c2, _, err := dialer.Dial("ws"+strings.TrimPrefix(srv2.URL, "http")+"/ws", nil) if err != nil { t.Fatalf("dial: %v", err) } defer c2.Close() if err := c2.WriteMessage(websocket.TextMessage, []byte("hi")); err != nil { t.Fatalf("write: %v", err) } deadline := time.Now().Add(2 * time.Second) var srvConn *ws.Connection for time.Now().Before(deadline) { if v := got.Load(); v != nil { srvConn = v.(*ws.Connection) break } time.Sleep(5 * time.Millisecond) } if srvConn == nil { t.Fatal("timeout") } hub.Register(srvConn) deadline = time.Now().Add(time.Second) for time.Now().Before(deadline) { if hub.Count() == 1 { break } time.Sleep(5 * time.Millisecond) } if hub.Count() != 1 { t.Fatalf("Hub Count = %d, want 1", hub.Count()) } // 直接关闭服务端连接(closeChan 关闭),使后续 Send 返回 error。 srvConn.Close() // 给 closeChan 关闭传播一点时间。 time.Sleep(20 * time.Millisecond) // 广播必须在超时内返回(不阻塞 Hub),且失败连接被行内移除。 done := make(chan struct{}) go func() { hub.Broadcast([]byte("should-not-deadlock")) close(done) }() select { case <-done: case <-time.After(3 * time.Second): t.Fatal("Hub.Broadcast deadlocked (C2a: unregister self-send loop)") } // 等待 Hub 内部清理失败的连接。 deadline = time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { if hub.Count() == 0 { break } time.Sleep(10 * time.Millisecond) } if hub.Count() != 0 { t.Errorf("Hub Count = %d, want 0 (failed conn should be removed inline)", hub.Count()) } } // 回归 C2a/C2b:Hub 广播到正常连接,客户端能收到。 func TestHubBroadcastReachesClients(t *testing.T) { hub := ws.NewHub() go hub.Run() // 用带 Hub 注册的 handler。 handler := &ws.DefaultHandler{ OnConnectFunc: func(conn *ws.Connection) { hub.Register(conn) }, } srv, dialer := newWSTestServer(t, handler) // 两个客户端。 c1 := dial(t, srv, dialer) c2 := dial(t, srv, dialer) defer c1.Close() defer c2.Close() // 等待两个连接注册。 deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) && hub.Count() < 2 { time.Sleep(10 * time.Millisecond) } if hub.Count() != 2 { t.Fatalf("Hub Count = %d, want 2", hub.Count()) } // 广播 JSON。 if err := hub.BroadcastJSON(map[string]string{"msg": "hello"}); err != nil { t.Fatalf("BroadcastJSON: %v", err) } // 两个客户端都应收到。 for i, c := range []*websocket.Conn{c1, c2} { c.SetReadDeadline(time.Now().Add(2 * time.Second)) _, msg, err := c.ReadMessage() if err != nil { t.Fatalf("client %d read: %v", i, err) } var m map[string]string if err := json.Unmarshal(msg, &m); err != nil { t.Fatalf("client %d unmarshal %q: %v", i, string(msg), err) } if m["msg"] != "hello" { t.Errorf("client %d got %q, want hello", i, m["msg"]) } } } // 回归 C2a-residual:Send 非阻塞——缓冲满立即返回 ErrSendBufferFull 而非阻塞等待。 // 旧实现(阻塞 select)在 writePump 退出但 closeChan 未关、send 缓冲满时会阻塞最长 pongWait, // 导致 Hub 广播持写锁期间 stall。非阻塞投递保证持锁期间永不阻塞。 // 用 internal test(ws_send_internal_test.go)直接测私有 send channel,此处仅验证公开行为: // 连接 Close 后,Send 立即返回(不阻塞)。 func TestSendNonBlockingOnClosed(t *testing.T) { var got atomic.Value r := gin.New() r.GET("/ws", ws.HandleFunc(func(conn *ws.Connection, message []byte) { got.Store(conn) })) srv := httptest.NewServer(r) t.Cleanup(srv.Close) dialer := &websocket.Dialer{HandshakeTimeout: 2 * time.Second} c, _, err := dialer.Dial("ws"+strings.TrimPrefix(srv.URL, "http")+"/ws", nil) if err != nil { t.Fatalf("dial: %v", err) } defer c.Close() if err := c.WriteMessage(websocket.TextMessage, []byte("hi")); err != nil { t.Fatalf("write: %v", err) } deadline := time.Now().Add(2 * time.Second) var srvConn *ws.Connection for time.Now().Before(deadline) { if v := got.Load(); v != nil { srvConn = v.(*ws.Connection) break } time.Sleep(5 * time.Millisecond) } if srvConn == nil { t.Fatal("timeout") } srvConn.Close() // Close 后连续 Send 必须立即返回 error,不阻塞。 done := make(chan struct{}) go func() { defer close(done) for i := 0; i < 1000; i++ { _ = srvConn.Send([]byte("x")) } }() select { case <-done: case <-time.After(3 * time.Second): t.Fatal("Send after Close blocked (should be non-blocking)") } } // ===== C2c:半开连接超时退出(不永久泄漏 goroutine)===== // 回归 C2c:客户端断开后,服务端读循环应在 pongWait 超时内退出(非永久阻塞)。 // 用短超时配置不便(常量包级),此处用真实连接断开 + 有限等待验证 OnClose 被调用。 func TestHalfOpenConnectionExitsOnClose(t *testing.T) { var closed atomic.Int32 handler := &ws.DefaultHandler{ OnCloseFunc: func(conn *ws.Connection) { closed.Add(1) }, } srv, dialer := newWSTestServer(t, handler) c := dial(t, srv, dialer) // 客户端主动关闭(发 close 帧)。 c.Close() // 服务端应在有限时间内 OnClose。pongWait=60s 太长,但客户端发的是正常 close, // ReadMessage 立即返回 close 错误,OnClose 应很快触发。 deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { if closed.Load() > 0 { return } time.Sleep(10 * time.Millisecond) } t.Errorf("OnClose not called within timeout (C2c: read loop should exit on client close), closed=%d", closed.Load()) } // ===== 辅助:验证 errors 导入被使用 ===== var _ = errors.New var _ = http.StatusOK // TestHubStopConcurrentNoPanic H-8 回归:并发 Stop 不应 double-close panic。 // 修复前 Stop 用 select{<-stop/default:close(stop)},两个 goroutine 都走 default // 同时 close → panic。修复后 stopOnce 保证 close 仅一次。 func TestHubStopConcurrentNoPanic(t *testing.T) { hub := ws.NewHub() go hub.Run() var wg sync.WaitGroup for i := 0; i < 10; i++ { wg.Add(1) go func() { defer wg.Done() hub.Stop() // 并发 Stop }() } wg.Wait() // 再调一次 Stop(已 stop)也应安全返回 hub.Stop() } // TestHubStopBeforeRunNoPanic H-9 回归:Stop 先于 Run 调用,随后 Run 不应 // 触发 wg 负计数 panic。 func TestHubStopBeforeRunNoPanic(t *testing.T) { hub := ws.NewHub() hub.Stop() // Run 尚未启动,Wait 立即返回 // 之后启动 Run(实际场景是误用),应安全退出而非 panic go hub.Run() // 给 Run 一点时间观察到 stop 已 close 并退出 time.Sleep(50 * time.Millisecond) // 再次 Stop 确保 wg 归零(Run 的 runOnce 已执行,wg.Add/Done 配对) hub.Stop() }