03f379c73c
- Add client struct with conn and room fields
- Add broadcastMsg struct with room and data fields
- Change Hub.clients to Hub.rooms map[string]map[*websocket.Conn]bool
- Add Hub.connRoom map[*websocket.Conn]string for reverse lookup
- Change broadcast channel type to chan broadcastMsg
- Change register channel type to chan client
- Update New() to initialize rooms and connRoom maps
- Update ClientCount() to use len(h.connRoom)
- Add RoomCount() method
- Update Run() loop for room-segmented register/unregister/broadcast
- Update HandleWebSocket to extract room from query param
- Backward compatible: clients without ?room use default empty room
- Update TestNew to verify rooms and connRoom maps initialized
- Add TestRoomCount to verify initial room count is 0
- Fix TestBroadcastChannel to use broadcastMsg type
All existing unit and integration tests pass (16 hub tests + 21 other).
🤖 Assisted by the code-assist SOP
88 lines
1.6 KiB
Go
88 lines
1.6 KiB
Go
package hub
|
|
|
|
import (
|
|
"bytes"
|
|
"testing"
|
|
"time"
|
|
|
|
"websocket-relay/internal/logging"
|
|
)
|
|
|
|
func newTestLogger() *logging.Logger {
|
|
return logging.NewLogger("debug", &bytes.Buffer{})
|
|
}
|
|
|
|
func TestNew(t *testing.T) {
|
|
h := New(newTestLogger())
|
|
if h == nil {
|
|
t.Fatal("New returned nil")
|
|
}
|
|
if h.rooms == nil {
|
|
t.Error("rooms map not initialized")
|
|
}
|
|
if h.connRoom == nil {
|
|
t.Error("connRoom map not initialized")
|
|
}
|
|
if h.broadcast == nil {
|
|
t.Error("broadcast channel not initialized")
|
|
}
|
|
if h.stop == nil {
|
|
t.Error("stop channel not initialized")
|
|
}
|
|
}
|
|
|
|
func TestClientCount(t *testing.T) {
|
|
h := New(newTestLogger())
|
|
go h.Run()
|
|
defer h.Shutdown()
|
|
|
|
if count := h.ClientCount(); count != 0 {
|
|
t.Errorf("Expected 0 clients, got %d", count)
|
|
}
|
|
}
|
|
|
|
func TestRoomCount(t *testing.T) {
|
|
h := New(newTestLogger())
|
|
go h.Run()
|
|
defer h.Shutdown()
|
|
|
|
if count := h.RoomCount(); count != 0 {
|
|
t.Errorf("Expected 0 rooms, got %d", count)
|
|
}
|
|
}
|
|
|
|
func TestBroadcastChannel(t *testing.T) {
|
|
h := New(newTestLogger())
|
|
go h.Run()
|
|
defer h.Shutdown()
|
|
|
|
select {
|
|
case h.broadcast <- broadcastMsg{room: "", data: []byte("test")}:
|
|
// Channel is working
|
|
case <-time.After(100 * time.Millisecond):
|
|
t.Error("broadcast channel blocked")
|
|
}
|
|
}
|
|
|
|
func TestShutdown(t *testing.T) {
|
|
h := New(newTestLogger())
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
h.Run()
|
|
close(done)
|
|
}()
|
|
|
|
// Ensure Run is processing before shutdown
|
|
time.Sleep(10 * time.Millisecond)
|
|
|
|
h.Shutdown()
|
|
|
|
select {
|
|
case <-done:
|
|
// Hub.Run() returned successfully
|
|
case <-time.After(1 * time.Second):
|
|
t.Fatal("Hub.Run() did not return after Shutdown")
|
|
}
|
|
}
|