Messenger/server/internal/realtime/hub_test.go

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