- Updated the WebSocket hub to replace the insecure `user_id` query parameter with an authentication token for secure connections. - Introduced typed events for mail operations (created, updated, deleted) to streamline event handling. - Implemented heartbeat functionality (ping/pong) to maintain connection health. - Enhanced client reconnection logic and delta replay for improved user experience. - Added limits on connections per user/session to prevent abuse and ensure stability.
99 lines
2.8 KiB
Go
99 lines
2.8 KiB
Go
package realtime
|
|
|
|
import (
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestExtractToken(t *testing.T) {
|
|
req := httptest.NewRequest("GET", "/ws", nil)
|
|
req.Header.Set("Authorization", "Bearer abc")
|
|
token, ok := extractToken(req)
|
|
if !ok || token != "abc" {
|
|
t.Fatalf("extractToken auth header = (%q, %v), want (abc, true)", token, ok)
|
|
}
|
|
|
|
req = httptest.NewRequest("GET", "/ws?token=q1", nil)
|
|
token, ok = extractToken(req)
|
|
if !ok || token != "q1" {
|
|
t.Fatalf("extractToken token query = (%q, %v), want (q1, true)", token, ok)
|
|
}
|
|
|
|
req = httptest.NewRequest("GET", "/ws?access_token=q2", nil)
|
|
token, ok = extractToken(req)
|
|
if !ok || token != "q2" {
|
|
t.Fatalf("extractToken access_token query = (%q, %v), want (q2, true)", token, ok)
|
|
}
|
|
}
|
|
|
|
func TestParseSinceCursor(t *testing.T) {
|
|
req := httptest.NewRequest("GET", "/ws?since=42", nil)
|
|
if got := parseSinceCursor(req); got != 42 {
|
|
t.Fatalf("parseSinceCursor() = %d, want 42", got)
|
|
}
|
|
|
|
req = httptest.NewRequest("GET", "/ws?since=bad", nil)
|
|
if got := parseSinceCursor(req); got != 0 {
|
|
t.Fatalf("parseSinceCursor() invalid = %d, want 0", got)
|
|
}
|
|
}
|
|
|
|
func TestTokenSessionKeyStable(t *testing.T) {
|
|
a := tokenSessionKey("same-token")
|
|
b := tokenSessionKey("same-token")
|
|
c := tokenSessionKey("other-token")
|
|
if a != b {
|
|
t.Fatalf("session keys for same token differ: %q != %q", a, b)
|
|
}
|
|
if a == c {
|
|
t.Fatalf("session keys for different tokens should differ")
|
|
}
|
|
}
|
|
|
|
func TestRegisterRespectsLimits(t *testing.T) {
|
|
h := NewHub(nil, nil)
|
|
h.SetLimits(2, 1, 10)
|
|
|
|
c1 := &conn{userID: "u1", sessionKey: "s1"}
|
|
c2 := &conn{userID: "u1", sessionKey: "s2"}
|
|
c3 := &conn{userID: "u1", sessionKey: "s3"}
|
|
|
|
if err := h.register(c1); err != nil {
|
|
t.Fatalf("register c1 error = %v", err)
|
|
}
|
|
if err := h.register(c2); err != nil {
|
|
t.Fatalf("register c2 error = %v", err)
|
|
}
|
|
if err := h.register(c3); err == nil {
|
|
t.Fatalf("register c3 expected user limit error")
|
|
}
|
|
|
|
h2 := NewHub(nil, nil)
|
|
h2.SetLimits(5, 1, 10)
|
|
if err := h2.register(&conn{userID: "u1", sessionKey: "s1"}); err != nil {
|
|
t.Fatalf("register first session conn error = %v", err)
|
|
}
|
|
if err := h2.register(&conn{userID: "u1", sessionKey: "s1"}); err == nil {
|
|
t.Fatalf("register second session conn expected session limit error")
|
|
}
|
|
}
|
|
|
|
func TestBroadcastAssignsSequenceAndKeepsReplayBuffer(t *testing.T) {
|
|
h := NewHub(nil, nil)
|
|
h.SetLimits(10, 10, 2)
|
|
|
|
h.Broadcast("u1", NewMailCreatedEvent("m1", "a1"))
|
|
h.Broadcast("u1", NewMailUpdatedEvent("m2", "a1"))
|
|
h.Broadcast("u1", NewMailDeletedEvent("m3", "a1"))
|
|
|
|
if head := h.HistoryHead("u1"); head != 3 {
|
|
t.Fatalf("HistoryHead() = %d, want 3", head)
|
|
}
|
|
if len(h.history["u1"]) != 2 {
|
|
t.Fatalf("history length = %d, want 2", len(h.history["u1"]))
|
|
}
|
|
if h.history["u1"][0].Seq != 2 || h.history["u1"][1].Seq != 3 {
|
|
t.Fatalf("history seqs = [%d, %d], want [2, 3]", h.history["u1"][0].Seq, h.history["u1"][1].Seq)
|
|
}
|
|
}
|