kenn-io--agentsview/internal/server/timeout_custom_test.go

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))
		})
	}
}