Add automated test to chat example

Author: nhooyr

chat-example/chat.go

@@ -15,8 +15,28 @@ import (
 
 // chatServer enables broadcasting to a set of subscribers.
 type chatServer struct {
+	registerOnce sync.Once
+	m            http.ServeMux
+
 	subscribersMu sync.RWMutex
-	subscribers   map[chan<- []byte]struct{}
+	subscribers   map[*subscriber]struct{}
+}
+
+// subscriber represents a subscriber.
+// Messages are sent on the msgs channel and if the client
+// cannot keep up with the messages, closeSlow is called.
+type subscriber struct {
+	msgs      chan []byte
+	closeSlow func()
+}
+
+func (cs *chatServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	cs.registerOnce.Do(func() {
+		cs.m.Handle("/", http.FileServer(http.Dir(".")))
+		cs.m.HandleFunc("/subscribe", cs.subscribeHandler)
+		cs.m.HandleFunc("/publish", cs.publishHandler)
+	})
+	cs.m.ServeHTTP(w, r)
 }
 
 // subscribeHandler accepts the WebSocket connection and then subscribes
@@ -57,11 +77,13 @@ func (cs *chatServer) publishHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	cs.publish(msg)
+
+	w.WriteHeader(http.StatusAccepted)
 }
 
 // subscribe subscribes the given WebSocket to all broadcast messages.
-// It creates a msgs chan with a buffer of 16 to give some room to slower
-// connections and then registers it. It then listens for all messages
+// It creates a subscriber with a buffered msgs chan to give some room to slower
+// connections and then registers the subscriber. It then listens for all messages
 // and writes them to the WebSocket. If the context is cancelled or
 // an error occurs, it returns and deletes the subscription.
 //
@@ -70,13 +92,18 @@ func (cs *chatServer) publishHandler(w http.ResponseWriter, r *http.Request) {
 func (cs *chatServer) subscribe(ctx context.Context, c *websocket.Conn) error {
 	ctx = c.CloseRead(ctx)
 
-	msgs := make(chan []byte, 16)
-	cs.addSubscriber(msgs)
-	defer cs.deleteSubscriber(msgs)
+	s := &subscriber{
+		msgs: make(chan []byte, 16),
+		closeSlow: func() {
+			c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages")
+		},
+	}
+	cs.addSubscriber(s)
+	defer cs.deleteSubscriber(s)
 
 	for {
 		select {
-		case msg := <-msgs:
+		case msg := <-s.msgs:
 			err := writeTimeout(ctx, time.Second*5, c, msg)
 			if err != nil {
 				return err
@@ -94,29 +121,29 @@ func (cs *chatServer) publish(msg []byte) {
 	cs.subscribersMu.RLock()
 	defer cs.subscribersMu.RUnlock()
 
-	for c := range cs.subscribers {
+	for s := range cs.subscribers {
 		select {
-		case c <- msg:
+		case s.msgs <- msg:
 		default:
+			go s.closeSlow()
 		}
 	}
 }
 
-// addSubscriber registers a subscriber with a channel
-// on which to send messages.
-func (cs *chatServer) addSubscriber(msgs chan<- []byte) {
+// addSubscriber registers a subscriber.
+func (cs *chatServer) addSubscriber(s *subscriber) {
 	cs.subscribersMu.Lock()
 	if cs.subscribers == nil {
-		cs.subscribers = make(map[chan<- []byte]struct{})
+		cs.subscribers = make(map[*subscriber]struct{})
 	}
-	cs.subscribers[msgs] = struct{}{}
+	cs.subscribers[s] = struct{}{}
 	cs.subscribersMu.Unlock()
 }
 
-// deleteSubscriber deletes the subscriber with the given msgs channel.
-func (cs *chatServer) deleteSubscriber(msgs chan []byte) {
+// deleteSubscriber deletes the given subscriber.
+func (cs *chatServer) deleteSubscriber(s *subscriber) {
 	cs.subscribersMu.Lock()
-	delete(cs.subscribers, msgs)
+	delete(cs.subscribers, s)
 	cs.subscribersMu.Unlock()
 }
 

chat-example/chat_test.go

@@ -0,0 +1,137 @@
+// +build !js
+
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"nhooyr.io/websocket"
+)
+
+func TestGrace(t *testing.T) {
+	t.Parallel()
+
+	var cs chatServer
+	var g websocket.Grace
+	s := httptest.NewServer(g.Handler(&cs))
+	defer s.Close()
+	defer g.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	cl1, err := newClient(ctx, s.URL)
+	assertSuccess(t, err)
+	defer cl1.Close()
+
+	cl2, err := newClient(ctx, s.URL)
+	assertSuccess(t, err)
+	defer cl2.Close()
+
+	err = cl1.publish(ctx, "hello")
+	assertSuccess(t, err)
+
+	assertReceivedMessage(ctx, cl1, "hello")
+	assertReceivedMessage(ctx, cl2, "hello")
+}
+
+type client struct {
+	msgs chan string
+	url  string
+	c    *websocket.Conn
+}
+
+func newClient(ctx context.Context, url string) (*client, error) {
+	wsURL := strings.ReplaceAll(url, "http://", "ws://")
+	c, _, err := websocket.Dial(ctx, wsURL+"/subscribe", nil)
+	if err != nil {
+		return nil, err
+	}
+
+	cl := &client{
+		msgs: make(chan string, 16),
+		url:  url,
+		c:    c,
+	}
+	go cl.readLoop()
+
+	return cl, nil
+}
+
+func (cl *client) readLoop() {
+	defer cl.c.Close(websocket.StatusInternalError, "")
+	defer close(cl.msgs)
+
+	for {
+		typ, b, err := cl.c.Read(context.Background())
+		if err != nil {
+			return
+		}
+
+		if typ != websocket.MessageText {
+			cl.c.Close(websocket.StatusUnsupportedData, "expected text message")
+			return
+		}
+
+		select {
+		case cl.msgs <- string(b):
+		default:
+			cl.c.Close(websocket.StatusInternalError, "messages coming in too fast to handle")
+			return
+		}
+	}
+}
+
+func (cl *client) receive(ctx context.Context) (string, error) {
+	select {
+	case msg, ok := <-cl.msgs:
+		if !ok {
+			return "", errors.New("client closed")
+		}
+		return msg, nil
+	case <-ctx.Done():
+		return "", ctx.Err()
+	}
+}
+
+func (cl *client) publish(ctx context.Context, msg string) error {
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, cl.url+"/publish", strings.NewReader(msg))
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusAccepted {
+		return fmt.Errorf("publish request failed: %v", resp.StatusCode)
+	}
+	return nil
+}
+
+func (cl *client) Close() error {
+	return cl.c.Close(websocket.StatusNormalClosure, "")
+}
+
+func assertSuccess(t *testing.T, err error) {
+	t.Helper()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func assertReceivedMessage(ctx context.Context, cl *client, msg string) error {
+	msg, err := cl.receive(ctx)
+	if err != nil {
+		return err
+	}
+	if msg != "hello" {
+		return fmt.Errorf("expected hello but got %q", msg)
+	}
+	return nil
+}

chat-example/go.mod

@@ -35,16 +35,10 @@ func run() error {
 	}
 	log.Printf("listening on http://%v", l.Addr())
 
-	var ws chatServer
-
-	m := http.NewServeMux()
-	m.Handle("/", http.FileServer(http.Dir(".")))
-	m.HandleFunc("/subscribe", ws.subscribeHandler)
-	m.HandleFunc("/publish", ws.publishHandler)
-
+	var cs chatServer
 	var g websocket.Grace
 	s := http.Server{
-		Handler:      g.Handler(m),
+		Handler:      g.Handler(&cs),
 		ReadTimeout:  time.Second * 10,
 		WriteTimeout: time.Second * 10,
 	}