More coverage for compress middleware

Signed-off-by: Vishal Rana <[email protected]>

echo.go

@@ -8,14 +8,12 @@ import (
 	"io"
 	"log"
 	"net/http"
-	"path"
+	spath "path"
 	"reflect"
 	"runtime"
 	"strings"
 	"sync"
 
-	"path/filepath"
-
 	"github.com/bradfitz/http2"
 	"github.com/mattn/go-colorable"
 	"golang.org/x/net/websocket"
@@ -118,6 +116,8 @@ const (
 	//-----------
 
 	WebSocket = "websocket"
+
+	indexFile = "index.html"
 )
 
 var (
@@ -326,7 +326,7 @@ func (e *Echo) ServeDir(path, dir string) {
 // ServeFile serves a file.
 func (e *Echo) ServeFile(path, file string) {
 	e.Get(path, func(c *Context) error {
-		dir, file := filepath.Split(file)
+		dir, file := spath.Split(file)
 		return serveFile(dir, file, c)
 	})
 }
@@ -340,7 +340,7 @@ func serveFile(dir, file string, c *Context) error {
 
 	fi, _ := f.Stat()
 	if fi.IsDir() {
-		file = path.Join(file, "index.html")
+		file = spath.Join(file, indexFile)
 		f, err = fs.Open(file)
 		if err != nil {
 			return NewHTTPError(http.StatusForbidden)

middleware/compress.go

@@ -25,8 +25,8 @@ func (w gzipWriter) Write(b []byte) (int, error) {
 	return w.Writer.Write(b)
 }
 
-func (w gzipWriter) Flush() {
-	w.Writer.(*gzip.Writer).Flush()
+func (w gzipWriter) Flush() error {
+	return w.Writer.(*gzip.Writer).Flush()
 }
 
 func (w gzipWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {

middleware/compress_test.go

@@ -6,17 +6,38 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/labstack/echo"
 	"github.com/stretchr/testify/assert"
 )
 
+type closeNotifyingRecorder struct {
+	*httptest.ResponseRecorder
+	closed chan bool
+}
+
+func newCloseNotifyingRecorder() *closeNotifyingRecorder {
+	return &closeNotifyingRecorder{
+		httptest.NewRecorder(),
+		make(chan bool, 1),
+	}
+}
+
+func (c *closeNotifyingRecorder) close() {
+	c.closed <- true
+}
+
+func (c *closeNotifyingRecorder) CloseNotify() <-chan bool {
+	return c.closed
+}
+
 func TestGzip(t *testing.T) {
 	req, _ := http.NewRequest(echo.GET, "/", nil)
 	rec := httptest.NewRecorder()
 	c := echo.NewContext(req, echo.NewResponse(rec), echo.New())
 	h := func(c *echo.Context) error {
-		c.Response().Write([]byte("test")) // Content-Type sniffing
+		c.Response().Write([]byte("test")) // For Content-Type sniffing
 		return nil
 	}
 
@@ -43,3 +64,58 @@ func TestGzip(t *testing.T) {
 		assert.Equal(t, "test", buf.String())
 	}
 }
+
+func TestGzipFlush(t *testing.T) {
+	rec := httptest.NewRecorder()
+	buf := new(bytes.Buffer)
+	w := gzip.NewWriter(buf)
+	gw := gzipWriter{Writer: w, ResponseWriter: rec}
+
+	n0 := buf.Len()
+	if n0 != 0 {
+		t.Fatalf("buffer size = %d before writes; want 0", n0)
+	}
+
+	if err := gw.Flush(); err != nil {
+		t.Fatal(err)
+	}
+
+	n1 := buf.Len()
+	if n1 == 0 {
+		t.Fatal("no data after first flush")
+	}
+
+	gw.Write([]byte("x"))
+
+	n2 := buf.Len()
+	if n1 != n2 {
+		t.Fatalf("after writing a single byte, size changed from %d to %d; want no change", n1, n2)
+	}
+
+	if err := gw.Flush(); err != nil {
+		t.Fatal(err)
+	}
+
+	n3 := buf.Len()
+	if n2 == n3 {
+		t.Fatal("Flush didn't flush any data")
+	}
+}
+
+func TestGzipCloseNotify(t *testing.T) {
+	rec := newCloseNotifyingRecorder()
+	buf := new(bytes.Buffer)
+	w := gzip.NewWriter(buf)
+	gw := gzipWriter{Writer: w, ResponseWriter: rec}
+	closed := false
+	notifier := gw.CloseNotify()
+	rec.close()
+
+	select {
+	case <-notifier:
+		closed = true
+	case <-time.After(time.Second):
+	}
+
+	assert.Equal(t, closed, true)
+}