52 lines
1.4 KiB
Go
52 lines
1.4 KiB
Go
package server
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
type observedDeadlineWriter struct {
|
|
*httptest.ResponseRecorder
|
|
writes, reads int
|
|
deadline time.Time
|
|
}
|
|
|
|
func (w *observedDeadlineWriter) SetWriteDeadline(d time.Time) error {
|
|
w.writes++
|
|
w.deadline = d
|
|
return nil
|
|
}
|
|
func (w *observedDeadlineWriter) SetReadDeadline(d time.Time) error { w.reads++; return nil }
|
|
|
|
func TestResponseWriteDeadlines(t *testing.T) {
|
|
raw := &observedDeadlineWriter{ResponseRecorder: httptest.NewRecorder()}
|
|
w := &deadlineWriter{ResponseWriter: raw}
|
|
w.WriteHeader(200)
|
|
if raw.writes != 1 || !raw.deadline.After(time.Now()) {
|
|
t.Fatal("headers lack deadline")
|
|
}
|
|
if _, err := w.Write([]byte("body")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := w.FlushError(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if raw.writes != 3 {
|
|
t.Fatal("write/flush must refresh deadline", raw.writes)
|
|
}
|
|
if err := http.NewResponseController(w).SetReadDeadline(time.Time{}); err != nil || raw.reads != 1 {
|
|
t.Fatal("wrapper must preserve read deadline control", err)
|
|
}
|
|
}
|
|
|
|
func TestNonStreamResponseHasWriteDeadline(t *testing.T) {
|
|
s := testServer(t)
|
|
raw := &observedDeadlineWriter{ResponseRecorder: httptest.NewRecorder()}
|
|
s.ServeHTTP(raw, httptest.NewRequest("GET", "/admin/api/session", nil))
|
|
if raw.Code != 200 || raw.writes == 0 || !raw.deadline.After(time.Now()) {
|
|
t.Fatal("nonstream response lacks write deadline")
|
|
}
|
|
}
|