increase test coverage for interceptor

Author: romainmenke

Diff for: http/interceptor_test.go

@@ -47,6 +47,75 @@ func TestWriteHeader(t *testing.T) {
 	}
 }
 
+func TestWrite(t *testing.T) {
+	waf, err := coraza.NewWAF(coraza.NewWAFConfig())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tx := waf.NewTransaction()
+	req, _ := http.NewRequest("GET", "", nil)
+	res := httptest.NewRecorder()
+
+	rw, responseProcessor := wrap(res, req, tx)
+	_, err = rw.Write([]byte("hello"))
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	_, err = rw.Write([]byte("world"))
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	err = responseProcessor(tx, req)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	if want, have := 200, res.Code; want != have {
+		t.Errorf("unexpected status code, want %d, have %d", want, have)
+	}
+}
+
+func TestWriteWithWriteHeader(t *testing.T) {
+	waf, err := coraza.NewWAF(coraza.NewWAFConfig())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tx := waf.NewTransaction()
+	req, _ := http.NewRequest("GET", "", nil)
+	res := httptest.NewRecorder()
+
+	rw, responseProcessor := wrap(res, req, tx)
+	rw.WriteHeader(204)
+	// although we called WriteHeader, status code should be applied until
+	// responseProcessor is called.
+	if unwanted, have := 204, res.Code; unwanted == have {
+		t.Errorf("unexpected status code %d", have)
+	}
+
+	_, err = rw.Write([]byte("hello"))
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	_, err = rw.Write([]byte("world"))
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	err = responseProcessor(tx, req)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	if want, have := 204, res.Code; want != have {
+		t.Errorf("unexpected status code, want %d, have %d", want, have)
+	}
+}
+
 func TestFlush(t *testing.T) {
 	waf, err := coraza.NewWAF(coraza.NewWAFConfig())
 	if err != nil {
@@ -75,6 +144,14 @@ func TestFlush(t *testing.T) {
 	}
 }
 
+type testReaderFrom struct {
+	io.Writer
+}
+
+func (x *testReaderFrom) ReadFrom(r io.Reader) (n int64, err error) {
+	return io.Copy(x, r)
+}
+
 func TestReadFrom(t *testing.T) {
 	waf, err := coraza.NewWAF(coraza.NewWAFConfig())
 	if err != nil {
@@ -84,13 +161,86 @@ func TestReadFrom(t *testing.T) {
 	tx := waf.NewTransaction()
 	req, _ := http.NewRequest("GET", "", nil)
 	res := httptest.NewRecorder()
-	rw, _ := wrap(res, req, tx)
+
+	type responseWriter interface {
+		http.ResponseWriter
+		http.Flusher
+	}
+
+	resWithReaderFrom := struct {
+		responseWriter
+		io.ReaderFrom
+	}{
+		res,
+		&testReaderFrom{res},
+	}
+
+	rw, responseProcessor := wrap(resWithReaderFrom, req, tx)
 	rw.WriteHeader(204)
+	// although we called WriteHeader, status code should be applied until
+	// responseProcessor is called.
+	if unwanted, have := 204, res.Code; unwanted == have {
+		t.Errorf("unexpected status code %d", have)
+	}
+
 	_, err = rw.(io.ReaderFrom).ReadFrom(bytes.NewBuffer([]byte("hello world")))
 	if err != nil {
 		t.Errorf("unexpected error: %v", err)
 	}
 
+	err = responseProcessor(tx, req)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	if want, have := 204, res.Code; want != have {
+		t.Errorf("unexpected status code, want %d, have %d", want, have)
+	}
+}
+
+type testPusher struct{}
+
+func (x *testPusher) Push(target string, opts *http.PushOptions) error {
+	return nil
+}
+
+func TestPusher(t *testing.T) {
+	waf, err := coraza.NewWAF(coraza.NewWAFConfig())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tx := waf.NewTransaction()
+	req, _ := http.NewRequest("GET", "", nil)
+	res := httptest.NewRecorder()
+
+	type responseWriter interface {
+		http.ResponseWriter
+		http.Flusher
+	}
+
+	resWithPush := struct {
+		responseWriter
+		http.Pusher
+	}{
+		res,
+		&testPusher{},
+	}
+
+	rw, responseProcessor := wrap(resWithPush, req, tx)
+	rw.WriteHeader(204)
+	rw.(http.Pusher).Push("http://example.com", nil)
+	// although we called WriteHeader, status code should be applied until
+	// responseProcessor is called.
+	if unwanted, have := 204, res.Code; unwanted == have {
+		t.Errorf("unexpected status code %d", have)
+	}
+
+	err = responseProcessor(tx, req)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
 	if want, have := 204, res.Code; want != have {
 		t.Errorf("unexpected status code, want %d, have %d", want, have)
 	}