package realtime import ( "context" "encoding/json" "net/http" "net/http/httptest" "strings" "sync" "testing" "time" "github.com/coder/websocket" ) type fakeMembers struct { mu sync.Mutex byUser map[int64][]int64 member map[[2]int64]bool lastReads [][2]int64 } func newFakeMembers() *fakeMembers { return &fakeMembers{ byUser: make(map[int64][]int64), member: make(map[[2]int64]bool), } } func (f *fakeMembers) ConversationIDsForUser(_ context.Context, userID int64) ([]int64, error) { f.mu.Lock() defer f.mu.Unlock() return f.byUser[userID], nil } func (f *fakeMembers) IsMember(_ context.Context, conversationID, userID int64) (bool, error) { f.mu.Lock() defer f.mu.Unlock() return f.member[[2]int64{conversationID, userID}], nil } func (f *fakeMembers) UpdateLastRead(_ context.Context, conversationID, userID int64) error { f.mu.Lock() defer f.mu.Unlock() f.lastReads = append(f.lastReads, [2]int64{conversationID, userID}) return nil } func (f *fakeMembers) readCount() int { f.mu.Lock() defer f.mu.Unlock() return len(f.lastReads) } func (f *fakeMembers) readAt(i int) [2]int64 { f.mu.Lock() defer f.mu.Unlock() return f.lastReads[i] } func waitFor(t *testing.T, cond func() bool, what string) { t.Helper() deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { if cond() { return } time.Sleep(10 * time.Millisecond) } t.Fatalf("timed out waiting for %s", what) } func startHub(t *testing.T, fm *fakeMembers) (*Hub, *httptest.Server) { t.Helper() hub := NewHub(fm) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { userID := int64(0) switch r.URL.Query().Get("user") { case "1": userID = 1 case "2": userID = 2 } hub.HandleWS(w, r, userID) })) t.Cleanup(srv.Close) return hub, srv } func dialWS(t *testing.T, srv *httptest.Server, user string) *websocket.Conn { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() url := "ws" + strings.TrimPrefix(srv.URL, "http") + "?user=" + user conn, _, err := websocket.Dial(ctx, url, nil) if err != nil { t.Fatalf("websocket dial: %v", err) } t.Cleanup(func() { conn.Close(websocket.StatusNormalClosure, "") }) time.Sleep(50 * time.Millisecond) return conn } func readEnvelope(t *testing.T, conn *websocket.Conn) envelope { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() _, data, err := conn.Read(ctx) if err != nil { t.Fatalf("ws read: %v", err) } var env envelope if err := json.Unmarshal(data, &env); err != nil { t.Fatalf("unmarshal envelope: %v", err) } return env } // readEventOfType reads events until one of the wanted type arrives, // skipping unrelated events such as presence. func readEventOfType(t *testing.T, conn *websocket.Conn, wantType string) envelope { t.Helper() deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { env := readEnvelope(t, conn) if env.Type == wantType { return env } } t.Fatalf("timed out waiting for event %q", wantType) return envelope{} } func sendEnvelope(t *testing.T, conn *websocket.Conn, env envelope) { t.Helper() data, err := json.Marshal(env) if err != nil { t.Fatalf("marshal envelope: %v", err) } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() if err := conn.Write(ctx, websocket.MessageText, data); err != nil { t.Fatalf("ws write: %v", err) } } func TestMessageCreatedReachesSubscribers(t *testing.T) { fm := newFakeMembers() fm.byUser[1] = []int64{10} fm.byUser[2] = []int64{10} fm.member[[2]int64{10, 1}] = true fm.member[[2]int64{10, 2}] = true hub, srv := startHub(t, fm) conn1 := dialWS(t, srv, "1") conn2 := dialWS(t, srv, "2") evt := MessageCreated{ConversationID: 10, MessageID: 55, SenderID: 1, Body: "hello"} hub.PublishMessageCreated(10, evt) for _, conn := range []*websocket.Conn{conn1, conn2} { env := readEventOfType(t, conn, EventMessageCreated) var got MessageCreated if err := json.Unmarshal(env.Data, &got); err != nil { t.Fatalf("unmarshal data: %v", err) } if got.MessageID != 55 || got.Body != "hello" { t.Errorf("data = %+v, want MessageID=55 Body=hello", got) } } } func TestNonMemberTypingIsIgnored(t *testing.T) { fm := newFakeMembers() fm.byUser[1] = []int64{10} fm.byUser[2] = []int64{10} fm.member[[2]int64{10, 2}] = true hub, srv := startHub(t, fm) conn1 := dialWS(t, srv, "1") conn2 := dialWS(t, srv, "2") sendEnvelope(t, conn1, envelope{ Type: EventTyping, Data: mustMarshal(clientTyping{ConversationID: 10, IsTyping: true}), }) evt := MessageCreated{ConversationID: 10, MessageID: 1, SenderID: 2, Body: "x"} hub.PublishMessageCreated(10, evt) env := readEventOfType(t, conn2, EventMessageCreated) if env.Type != EventMessageCreated { t.Errorf("member got %q, want message.created (no typing leak)", env.Type) } } func TestPresenceOnConnectAndClose(t *testing.T) { fm := newFakeMembers() fm.byUser[1] = []int64{10} fm.byUser[2] = []int64{10} _, srv := startHub(t, fm) conn1 := dialWS(t, srv, "1") conn2 := dialWS(t, srv, "2") env := readEnvelope(t, conn1) if env.Type != EventPresence { t.Fatalf("first event type = %q, want presence", env.Type) } var p Presence if err := json.Unmarshal(env.Data, &p); err != nil { t.Fatalf("unmarshal presence: %v", err) } if p.UserID != 2 || p.Status != StatusOnline { t.Errorf("presence = %+v, want user 2 online", p) } conn2.Close(websocket.StatusNormalClosure, "") env = readEnvelope(t, conn1) if env.Type != EventPresence { t.Fatalf("closing event type = %q, want presence", env.Type) } if err := json.Unmarshal(env.Data, &p); err != nil { t.Fatalf("unmarshal presence: %v", err) } if p.UserID != 2 || p.Status != StatusOffline { t.Errorf("presence = %+v, want user 2 offline", p) } } func TestReadEventUpdatesLastReadAndBroadcasts(t *testing.T) { fm := newFakeMembers() fm.byUser[1] = []int64{10} fm.byUser[2] = []int64{10} fm.member[[2]int64{10, 1}] = true fm.member[[2]int64{10, 2}] = true _, srv := startHub(t, fm) conn1 := dialWS(t, srv, "1") conn2 := dialWS(t, srv, "2") sendEnvelope(t, conn1, envelope{ Type: EventRead, Data: mustMarshal(clientRead{ConversationID: 10, MessageID: 99}), }) waitFor(t, func() bool { return fm.readCount() == 1 }, "UpdateLastRead call") if got := fm.readAt(0); got != [2]int64{10, 1} { t.Errorf("lastReads[0] = %v, want [10 1]", got) } env := readEnvelope(t, conn2) if env.Type != EventRead { t.Fatalf("type = %q, want read", env.Type) } var r Read if err := json.Unmarshal(env.Data, &r); err != nil { t.Fatalf("unmarshal read: %v", err) } if r.UserID != 1 || r.MessageID != 99 || r.ConversationID != 10 { t.Errorf("read = %+v, want user 1, message 99, conv 10", r) } }