Merge pull request #138 from melito/master

Add Flusher to the ResponseWriter wrapper

response_writer.go

@@ -12,6 +12,7 @@ import (
 // if the functionality calls for it.
 type ResponseWriter interface {
 	http.ResponseWriter
+	http.Flusher
 	// Status returns the status code of the response or 0 if the response has not been written.
 	Status() int
 	// Written returns whether or not the ResponseWriter has been written.
@@ -87,3 +88,10 @@ func (rw *responseWriter) callBefore() {
 		rw.beforeFuncs[i](rw)
 	}
 }
+
+func (rw *responseWriter) Flush() {
+	flusher, ok := rw.ResponseWriter.(http.Flusher)
+	if ok {
+		flusher.Flush()
+	}
+}

response_writer_test.go

@@ -2,6 +2,7 @@ package martini
 
 import (
 	"bufio"
+	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -140,3 +141,48 @@ func Test_ResponseWriter_CloseNotify(t *testing.T) {
 	}
 	expect(t, closed, true)
 }
+
+func Test_ResponseWriter_Flusher(t *testing.T) {
+
+	rec := httptest.NewRecorder()
+	rw := NewResponseWriter(rec)
+
+	_, ok := rw.(http.Flusher)
+	expect(t, ok, true)
+}
+
+func Test_ResponseWriter_FlusherHandler(t *testing.T) {
+
+	// New martini instance
+	m := Classic()
+
+	m.Get("/events", func(w http.ResponseWriter, r *http.Request) {
+
+		f, ok := w.(http.Flusher)
+		expect(t, ok, true)
+
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.Header().Set("Cache-Control", "no-cache")
+		w.Header().Set("Connection", "keep-alive")
+
+		for i := 0; i < 2; i++ {
+			time.Sleep(10 * time.Millisecond)
+			io.WriteString(w, "data: Hello\n\n")
+			f.Flush()
+		}
+
+	})
+
+	recorder := httptest.NewRecorder()
+	r, _ := http.NewRequest("GET", "/events", nil)
+	m.ServeHTTP(recorder, r)
+
+	if recorder.Code != 200 {
+		t.Error("Response not 200")
+	}
+
+	if recorder.Body.String() != "data: Hello\n\ndata: Hello\n\n" {
+		t.Error("Didn't receive correct body, got:", recorder.Body.String())
+	}
+
+}