Merge pull request #6 from launchdarkly/eb/ch81628/stream-deps

revise streaming HTTP helpers for cleaner implementation and fewer dependencies

Author: eli-darkly

Gopkg.lock

@@ -24,11 +24,6 @@
 #   go-tests = true
 #   unused-packages = true
 
-
-[[constraint]]
-  name = "github.com/launchdarkly/eventsource"
-  version = "1.2.0"
-
 [[constraint]]
   name = "github.com/stretchr/testify"
   version = "1.5.1"

Gopkg.toml

@@ -22,19 +22,6 @@ func TestClientFromHandler(t *testing.T) {
 }
 
 func TestClientFromHandlerConvertsPanicToError(t *testing.T) {
-	myError := errors.New("sorry")
-	handler := PanicHandler(myError)
-	client := ClientFromHandler(handler)
-
-	resp, err := client.Get("/")
-
-	expectedError := &url.Error{Op: "Get", URL: "/", Err: myError}
-	require.Error(t, err)
-	require.Nil(t, resp)
-	assert.Equal(t, expectedError, err)
-}
-
-func TestClientFromHandlerConvertsPanicWithArbitraryValueToError(t *testing.T) {
 	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		panic("sorry")
 	})

httphelpers/clients_test.go

@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"log"
 	"net/http"
+	"net/http/httptest"
 	"regexp"
 )
 
@@ -162,16 +163,28 @@ func SequentialHandler(firstHandler http.Handler, remainingHandlers ...http.Hand
 	})
 }
 
-// PanicHandler creates an HTTP handler that will panic.
+// BrokenConnectionHandler creates an HTTP handler that will simulate an I/O error.
 //
-// This is not useful inside an actual HTTP server, but when used in conjunction with NewHTTPClientFromHandler,
-// it allows you to simulate an I/O error.
+// When used with an httptest.Server, the handler forces an early close of the connection.
+// When used in a client created with ClientFromHandler, it causes a panic which is recovered
+// and converted to an error result. However, do not use this with httptest.ResponseRecorder
+// or your test will panic.
 //
-//     handler := PanicHandler(&net.AddrError{})
+//     handler := BrokenConnectionHandler()
 //     client := NewClientFromHandler(handler)
-//     // All requests made with this client will return an AddrError (or some other error that wraps that error)
-func PanicHandler(err error) http.Handler {
+//     // All requests made with this client will return an error
+func BrokenConnectionHandler() http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		panic(err)
+		if _, ok := w.(*httptest.ResponseRecorder); ok {
+			panic("httphelpers.BrokenConnectionHandler cannot be used with a ResponseRecorder")
+		}
+		if h, ok := w.(http.Hijacker); ok {
+			conn, _, err := h.Hijack()
+			if err == nil {
+				_ = conn.Close()
+				return
+			}
+		}
+		panic("connection deliberately closed by httphelpers.BrokenConnectionHandler")
 	})
 }

httphelpers/handlers.go

@@ -0,0 +1,108 @@
+package httphelpers
+
+import (
+	"bytes"
+	"fmt"
+	"net/http"
+)
+
+// SSEEvent is a simple representation of a Server-Sent Events message.
+type SSEEvent struct {
+	// ID is the optional unique ID of the event.
+	ID string
+
+	// Event is the message type, if any.
+	Event string
+
+	// Data is the message data.
+	Data string
+}
+
+// Bytes returns the stream data for the event.
+func (e SSEEvent) Bytes() []byte {
+	var buf bytes.Buffer
+	if e.ID != "" {
+		buf.WriteString(fmt.Sprintf("id: %s\n", e.ID))
+	}
+	if e.Event != "" {
+		buf.WriteString(fmt.Sprintf("event: %s\n", e.Event))
+	}
+	buf.WriteString(fmt.Sprintf("data: %s\n\n", e.Data))
+	return buf.Bytes()
+}
+
+// SSEStreamControl is the interface for manipulating streams created by SSEHandler.
+type SSEStreamControl interface {
+	// Enqueue is the same as Send, except that if there are currently no open connections to this
+	// endpoint, the event is enqueued and will be sent to the next client that connects.
+	Enqueue(event SSEEvent)
+
+	// Send sends an SSE event. If there are multiple open connections to this endpoint, the same
+	// event is sent to all of them. If there are no open connections, the event is discarded.
+	Send(event SSEEvent)
+
+	// EndAll terminates any existing connections to this endpoint, but allows new connections
+	// afterward.
+	EndAll()
+
+	// Close terminates any existing connections to this endpoint and causes the handler to reject any
+	// subsequent connections.
+	Close() error
+}
+
+type sseStreamControlImpl struct {
+	streamControl StreamControl
+}
+
+// SSEHandler creates an HTTP handler that streams Server-Sent Events data.
+//
+// The initialData parameter, if not nil, specifies a starting event that should always be sent to any
+// client that has connected to this endpoint. Then, any data provided via the SSEStreamControl interface
+// is copied to all connected clients. Connections remain open until either EndAll or Close is called on
+// on the SSEStreamControl.
+//
+// In this example, every request to this endpoint will receive an initial initEvent, and then another
+// event will be sent every second with a counter; every 30 seconds, all active stream connections are
+// are closed:
+//
+//     initialEvent := httphelpers.SSEEvent{Data: "hello"}
+//     handler, stream := httphelpers.SSEHandler(&initialEvent)
+//     (start server with handler)
+//     go func() {
+//         n := 1
+//         counter := time.NewTicker(time.Second)
+//         interrupter := time.NewTicker(time.Second * 10)
+//         for {
+//             select {
+//             case <-counter.C:
+//                 stream.Send(httphelpers.SSEEvent{Data: fmt.Sprintf("%d\n", n)})
+//                 n++
+//             case <-interrupter.C:
+//                 stream.EndAll()
+//             }
+//         }
+//     }()
+func SSEHandler(initialEvent *SSEEvent) (http.Handler, SSEStreamControl) {
+	var initialData []byte
+	if initialEvent != nil {
+		initialData = initialEvent.Bytes()
+	}
+	handler, streamControl := ChunkedStreamingHandler(initialData, "text/event-stream; charset=utf-8")
+	return handler, &sseStreamControlImpl{streamControl}
+}
+
+func (s *sseStreamControlImpl) Enqueue(event SSEEvent) {
+	s.streamControl.Enqueue(event.Bytes())
+}
+
+func (s *sseStreamControlImpl) Send(event SSEEvent) {
+	s.streamControl.Send(event.Bytes())
+}
+
+func (s *sseStreamControlImpl) EndAll() {
+	s.streamControl.EndAll()
+}
+
+func (s *sseStreamControlImpl) Close() error {
+	return s.streamControl.Close()
+}

httphelpers/handlers_sse.go

@@ -0,0 +1,47 @@
+package httphelpers
+
+import (
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSSEHandler(t *testing.T) {
+	initialEvent := SSEEvent{"id1", "event1", "data1"}
+	handler, stream := SSEHandler(&initialEvent)
+	defer stream.Close()
+
+	stream.Enqueue(SSEEvent{"", "event2", "data2"})
+	stream.Send(SSEEvent{"", "", "this isn't sent becauset here are no connections"})
+
+	WithServer(handler, func(server *httptest.Server) {
+		resp1, err := http.DefaultClient.Get(server.URL)
+		require.NoError(t, err)
+		defer resp1.Body.Close()
+
+		assert.Equal(t, 200, resp1.StatusCode)
+		assert.Equal(t, "text/event-stream; charset=utf-8", resp1.Header.Get("Content-Type"))
+
+		stream.Enqueue(SSEEvent{"", "event3", "data3"})
+		stream.EndAll()
+
+		data, err := ioutil.ReadAll(resp1.Body)
+
+		assert.NoError(t, err)
+		assert.Equal(t, `id: id1
+event: event1
+data: data1
+
+event: event2
+data: data2
+
+event: event3
+data: data3
+
+`, string(data))
+	})
+}