chore: refactor httpx.TimeoutHandler (#3400)

rest/handler/timeouthandler_test.go

@@ -17,35 +17,63 @@ import (
 )
 
 func TestTimeoutWriteFlushOutput(t *testing.T) {
-	timeoutHandler := TimeoutHandler(1000 * time.Millisecond)
-	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "text/event-stream;charset=utf-8")
-		flusher, ok := w.(http.Flusher)
-		if !ok {
-			http.Error(w, "Flushing not supported", http.StatusInternalServerError)
-			return
+	t.Run("flusher", func(t *testing.T) {
+		timeoutHandler := TimeoutHandler(1000 * time.Millisecond)
+		handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
+			flusher, ok := w.(http.Flusher)
+			if !ok {
+				http.Error(w, "Flushing not supported", http.StatusInternalServerError)
+				return
+			}
+
+			for i := 1; i <= 5; i++ {
+				fmt.Fprint(w, strconv.Itoa(i)+" cats\n\n")
+				flusher.Flush()
+				time.Sleep(time.Millisecond)
+			}
+		}))
+		req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
+		resp := httptest.NewRecorder()
+		handler.ServeHTTP(resp, req)
+		scanner := bufio.NewScanner(resp.Body)
+		var cats int
+		for scanner.Scan() {
+			line := scanner.Text()
+			if strings.Contains(line, "cats") {
+				cats++
+			}
 		}
-		for i := 1; i <= 5; i++ {
-			fmt.Fprint(w, strconv.Itoa(i)+"只猫猫\n\n")
-			flusher.Flush()
-			time.Sleep(time.Millisecond)
+		if err := scanner.Err(); err != nil {
+			cats = 0
 		}
-	}))
-	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
-	resp := httptest.NewRecorder()
-	handler.ServeHTTP(resp, req)
-	scanner := bufio.NewScanner(resp.Body)
-	mao := 0
-	for scanner.Scan() {
-		line := scanner.Text()
-		if strings.Contains(line, "猫猫") {
-			mao++
-		}
-	}
-	if err := scanner.Err(); err != nil {
-		mao = 0
-	}
-	assert.Equal(t, "5只猫猫", strconv.Itoa(mao)+"只猫猫")
+		assert.Equal(t, 5, cats)
+	})
+
+	t.Run("writer", func(t *testing.T) {
+		recorder := httptest.NewRecorder()
+		timeoutHandler := TimeoutHandler(1000 * time.Millisecond)
+		handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
+			flusher, ok := w.(http.Flusher)
+			if !ok {
+				http.Error(w, "Flushing not supported", http.StatusInternalServerError)
+				return
+			}
+
+			for i := 1; i <= 5; i++ {
+				fmt.Fprint(w, strconv.Itoa(i)+" cats\n\n")
+				flusher.Flush()
+				time.Sleep(time.Millisecond)
+				assert.Empty(t, recorder.Body.String())
+			}
+		}))
+		req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
+		resp := mockedResponseWriter{recorder}
+		handler.ServeHTTP(resp, req)
+		assert.Equal(t, "1 cats\n\n2 cats\n\n3 cats\n\n4 cats\n\n5 cats\n\n",
+			recorder.Body.String())
+	})
 }
 
 func TestTimeout(t *testing.T) {
@@ -274,3 +302,19 @@ func (m mockedPusher) WriteHeader(_ int) {
 func (m mockedPusher) Push(_ string, _ *http.PushOptions) error {
 	panic("implement me")
 }
+
+type mockedResponseWriter struct {
+	http.ResponseWriter
+}
+
+func (m mockedResponseWriter) Header() http.Header {
+	return m.ResponseWriter.Header()
+}
+
+func (m mockedResponseWriter) Write(bytes []byte) (int, error) {
+	return m.ResponseWriter.Write(bytes)
+}
+
+func (m mockedResponseWriter) WriteHeader(statusCode int) {
+	m.ResponseWriter.WriteHeader(statusCode)
+}