Repository: traefik/traefik
File: pkg/middlewares/compress/compress_test.go (Go)

Code before:
	handler, err := New(context.Background(), next, dynamic.Compress{}, "testing")
	require.NoError(t, err)

	rw := httptest.NewRecorder()
	handler.ServeHTTP(rw, req)

	assert.Equal(t, brotliValue, rw.Header().Get(contentEncodingHeader))
	assert.Equal(t, acceptEncodingHeader, rw.Header().Get(varyHeader))

	got, err := io.ReadAll(brotli.NewReader(rw.Body))
	require.NoError(t, err)
	assert.Equal(t, got, fakeBody)
}

func TestShouldNotCompressWhenIdentityAcceptEncodingHeader(t *testing.T) {
	req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil)
	req.Header.Set(acceptEncodingHeader, "identity")

	fakeBody := generateBytes(gzhttp.DefaultMinSize)
	next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		_, err := rw.Write(fakeBody)
		if err != nil {
			http.Error(rw, err.Error(), http.StatusInternalServerError)
		}
	})
	handler, err := New(context.Background(), next, dynamic.Compress{}, "testing")
	require.NoError(t, err)

	rw := httptest.NewRecorder()
	handler.ServeHTTP(rw, req)

	assert.Empty(t, rw.Header().Get(contentEncodingHeader))
	assert.Empty(t, rw.Header().Get(varyHeader))
	assert.EqualValues(t, rw.Body.Bytes(), fakeBody)
}

func TestShouldNotCompressWhenEmptyAcceptEncodingHeader(t *testing.T) {
	req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil)
	req.Header.Set(acceptEncodingHeader, "")

	fakeBody := generateBytes(gzhttp.DefaultMinSize)
	next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		_, err := rw.Write(fakeBody)
		if err != nil {
			http.Error(rw, err.Error(), http.StatusInternalServerError)
		}
	})
	handler, err := New(context.Background(), next, dynamic.Compress{}, "testing")
	require.NoError(t, err)

	rw := httptest.NewRecorder()

Review comment (suggestion):
```suggestion
		if r.Header.Get(acceptEncodingHeader) != "identity" {
			http.Error(rw, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			return
		}

		_, err := rw.Write(fakeBody)
```

Code after:
	handler, err := New(context.Background(), next, dynamic.Compress{}, "testing")
	require.NoError(t, err)

	rw := httptest.NewRecorder()
	handler.ServeHTTP(rw, req)

	assert.Equal(t, brotliValue, rw.Header().Get(contentEncodingHeader))
	assert.Equal(t, acceptEncodingHeader, rw.Header().Get(varyHeader))

	got, err := io.ReadAll(brotli.NewReader(rw.Body))
	require.NoError(t, err)
	assert.Equal(t, got, fakeBody)
}

func TestShouldNotCompressWhenIdentityAcceptEncodingHeader(t *testing.T) {
	req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil)
	req.Header.Set(acceptEncodingHeader, "identity")

	fakeBody := generateBytes(gzhttp.DefaultMinSize)
	next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		if r.Header.Get(acceptEncodingHeader) != "identity" {
			http.Error(rw, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			return
		}

		_, err := rw.Write(fakeBody)
		if err != nil {
			http.Error(rw, err.Error(), http.StatusInternalServerError)
		}
	})
	handler, err := New(context.Background(), next, dynamic.Compress{}, "testing")
	require.NoError(t, err)

	rw := httptest.NewRecorder()
	handler.ServeHTTP(rw, req)

	assert.Empty(t, rw.Header().Get(contentEncodingHeader))
	assert.Empty(t, rw.Header().Get(varyHeader))
	assert.EqualValues(t, rw.Body.Bytes(), fakeBody)
}

func TestShouldNotCompressWhenEmptyAcceptEncodingHeader(t *testing.T) {
	req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil)
	req.Header.Set(acceptEncodingHeader, "")

	fakeBody := generateBytes(gzhttp.DefaultMinSize)
	next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		if r.Header.Get(acceptEncodingHeader) != "" {
			http.Error(rw, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			return
		}