github.com/kyma-incubator/compass/components/director@v0.0.0-20230623144113-d764f56ff805/pkg/handler/timeout_handler_test.go (about)

     1  package handler_test
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"io"
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"sync"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/kyma-incubator/compass/components/director/pkg/correlation"
    14  	"github.com/kyma-incubator/compass/components/director/pkg/log"
    15  	"github.com/sirupsen/logrus"
    16  	"github.com/sirupsen/logrus/hooks/test"
    17  	"github.com/stretchr/testify/assert"
    18  
    19  	"github.com/kyma-incubator/compass/components/director/pkg/apperrors"
    20  	"github.com/kyma-incubator/compass/components/director/pkg/handler"
    21  	"github.com/stretchr/testify/require"
    22  )
    23  
    24  func TestHandlerWithTimeout_ReturnsTimeoutMessage(t *testing.T) {
    25  	timeout := time.Millisecond * 100
    26  	h, wait := getStubHandleFunc(t, timeout)
    27  	defer wait()
    28  
    29  	handlerWithTimeout, err := handler.WithTimeout(h, timeout)
    30  	require.NoError(t, err)
    31  
    32  	req := httptest.NewRequest(http.MethodGet, "/test", &bytes.Buffer{})
    33  	w := httptest.NewRecorder()
    34  
    35  	handlerWithTimeout.ServeHTTP(w, req)
    36  
    37  	resp := w.Result()
    38  	require.NotNil(t, resp)
    39  
    40  	require.Equal(t, handler.HeaderContentTypeValue, resp.Header.Get(handler.HeaderContentTypeKey))
    41  
    42  	respBody, err := io.ReadAll(resp.Body)
    43  	require.NoError(t, err)
    44  
    45  	actualError := getErrorMessage(t, respBody)
    46  
    47  	require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
    48  
    49  	require.Contains(t, actualError, "timed out")
    50  }
    51  
    52  func TestHandlerWithTimeout_LogCorrelationID(t *testing.T) {
    53  	timeout := time.Millisecond * 100
    54  	h, wait := getStubHandleFunc(t, timeout)
    55  	defer wait()
    56  
    57  	handlerWithTimeout, err := handler.WithTimeout(h, timeout)
    58  	require.NoError(t, err)
    59  
    60  	req := httptest.NewRequest(http.MethodGet, "/test", &bytes.Buffer{})
    61  	w := httptest.NewRecorder()
    62  
    63  	logger, hook := test.NewNullLogger()
    64  
    65  	ctx := log.ContextWithLogger(req.Context(), logrus.NewEntry(logger))
    66  	req = req.WithContext(ctx)
    67  
    68  	handlerWithTimeout.ServeHTTP(w, req)
    69  
    70  	reqID, ok := hook.LastEntry().Data[correlation.RequestIDHeaderKey]
    71  	require.True(t, ok)
    72  
    73  	assert.NotEqual(t, log.Configuration().BootstrapCorrelationID, reqID)
    74  }
    75  
    76  func getStubHandleFunc(t *testing.T, timeout time.Duration) (http.HandlerFunc, func()) {
    77  	wg := sync.WaitGroup{}
    78  	wg.Add(1)
    79  
    80  	sleepTime := timeout + (time.Millisecond * 10)
    81  
    82  	h := http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
    83  		defer wg.Done()
    84  		time.Sleep(sleepTime)
    85  		_, err := writer.Write([]byte("test"))
    86  		require.Equal(t, http.ErrHandlerTimeout, err)
    87  	})
    88  
    89  	wait := func() {
    90  		wg.Wait()
    91  	}
    92  
    93  	return h, wait
    94  }
    95  
    96  func getErrorMessage(t *testing.T, data []byte) string {
    97  	var body apperrors.Error
    98  	err := json.Unmarshal(data, &body)
    99  	require.NoError(t, err)
   100  	return body.Message
   101  }