Use require.NoError() to assert errors

msgbus_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"nhooyr.io/websocket"
 	"nhooyr.io/websocket/wsjson"
 )
@@ -52,7 +53,20 @@ func TestMessageBusPutGet(t *testing.T) {
 	assert.Equal(t, actual, expected)
 }
 
-func TestServeHTTPGETEmpty(t *testing.T) {
+func TestMessageBusSubscribe(t *testing.T) {
+	mb := New(nil)
+
+	msgs := mb.Subscribe("id1", "foo")
+
+	topic := mb.NewTopic("foo")
+	expected := Message{Topic: topic, Payload: []byte("foo")}
+	mb.Put(expected)
+
+	actual := <-msgs
+	assert.Equal(t, actual, expected)
+}
+
+func TestServeHTTPGETIndexEmpty(t *testing.T) {
 	assert := assert.New(t)
 
 	mb := New(nil)
@@ -157,6 +171,7 @@ func BenchmarkServeHTTPPOST(b *testing.B) {
 
 func TestServeHTTPSubscriber(t *testing.T) {
 	assert := assert.New(t)
+	require := require.New(t)
 
 	mb := New(nil)
 
@@ -171,16 +186,13 @@ func TestServeHTTPSubscriber(t *testing.T) {
 
 		// u := fmt.Sprintf("ws%s/hello", strings.TrimPrefix(s.URL, "http"))
 		ws, _, err := websocket.Dial(context.Background(), s.URL+"/hello", nil)
-
-		// ws, _, err := websocket.DefaultDialer.Dial(u, nil)
-		assert.NoError(err)
+		require.NoError(err)
 		defer ws.Close(websocket.StatusNormalClosure, "")
 
 		ready <- true
 
-		wsjson.Read(context.Background(), ws, &msg)
-		// err = ws.ReadJSON(&msg)
-		// ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+		err = wsjson.Read(context.Background(), ws, &msg)
+		require.NoError(err)
 
 		msgs <- msg
 	}
@@ -192,7 +204,50 @@ func TestServeHTTPSubscriber(t *testing.T) {
 	c := s.Client()
 	b := bytes.NewBufferString("hello world")
 	r, err := c.Post(s.URL+"/hello", "text/plain", b)
-	assert.NoError(err)
+	require.NoError(err)
+	defer r.Body.Close()
+
+	msg := <-msgs
+	assert.Equal(msg.ID, uint64(0))
+	assert.Equal(msg.Topic.Name, "hello")
+	assert.Equal(msg.Payload, []byte("hello world"))
+}
+
+func TestServeHTTPSubscriberReconnect(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	mb := New(nil)
+
+	s := httptest.NewServer(mb)
+	defer s.Close()
+
+	msgs := make(chan *Message)
+	ready := make(chan bool, 1)
+
+	consumer := func() {
+		var msg *Message
+
+		ws, _, err := websocket.Dial(context.Background(), s.URL+"/hello", nil)
+		require.NoError(err)
+		defer ws.Close(websocket.StatusNormalClosure, "")
+
+		ready <- true
+
+		err = wsjson.Read(context.Background(), ws, &msg)
+		require.NoError(err)
+
+		msgs <- msg
+	}
+
+	go consumer()
+
+	<-ready
+
+	c := s.Client()
+	b := bytes.NewBufferString("hello world")
+	r, err := c.Post(s.URL+"/hello", "text/plain", b)
+	require.NoError(err)
 	defer r.Body.Close()
 
 	msg := <-msgs