package server import ( "io" "net/http" "net/http/httptest" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestWithTimeout(t *testing.T) { t.Parallel() tests := []struct { name string timeout time.Duration handler http.HandlerFunc wantStatus int wantBody string wantHeaderKey string wantHeaderVal string assertResponse func(t *testing.T, resp *http.Response) }{ { name: "timeout", timeout: 10 * time.Millisecond, handler: func(w http.ResponseWriter, r *http.Request) { time.Sleep(50 * time.Millisecond) w.WriteHeader(http.StatusOK) w.Write([]byte("too slow")) }, assertResponse: assertTimeoutResponse, }, { name: "success", timeout: 100 * time.Millisecond, handler: func(w http.ResponseWriter, r *http.Request) { w.Header().Set("X-Custom", "value") w.WriteHeader(http.StatusCreated) w.Write([]byte(`{"status":"ok"}`)) }, wantStatus: http.StatusCreated, wantBody: `{"status":"ok"}`, wantHeaderKey: "X-Custom", wantHeaderVal: "value", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() s := newTestServerMinimal(t, tt.timeout) wrapped := s.withTimeout(tt.handler) req := httptest.NewRequest(http.MethodGet, "/", nil) w := httptest.NewRecorder() wrapped.ServeHTTP(w, req) resp := w.Result() defer resp.Body.Close() if tt.assertResponse != nil { tt.assertResponse(t, resp) return } assertRecorderStatus(t, w, tt.wantStatus) if tt.wantHeaderKey != "" { assert.Equal(t, tt.wantHeaderVal, resp.Header.Get(tt.wantHeaderKey), "header %s", tt.wantHeaderKey) } body, err := io.ReadAll(resp.Body) require.NoError(t, err) assert.Equal(t, tt.wantBody, string(body)) }) } }