upstream-mattermost/server/channels/web/response_writer_wrapper_test.go

// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
// See LICENSE.txt for license information.

package web

import (
	"bufio"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type TestHandler struct {
	TestFunc func(w http.ResponseWriter, r *http.Request)
}

func (h *TestHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.TestFunc(w, r)
}

type responseRecorderHijack struct {
	httptest.ResponseRecorder
}

func (r *responseRecorderHijack) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	r.WriteHeader(http.StatusOK)
	return nil, nil, nil
}

func newResponseWithHijack(original *httptest.ResponseRecorder) *responseRecorderHijack {
	return &responseRecorderHijack{*original}
}

func TestStatusCodeIsAccessible(t *testing.T) {
	resp := newWrappedWriter(httptest.NewRecorder())
	req := httptest.NewRequest("GET", "/api/v4/test", nil)
	handler := TestHandler{func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadRequest)
	}}
	handler.ServeHTTP(resp, req)
	assert.Equal(t, http.StatusBadRequest, resp.StatusCode())
}

func TestStatusCodeShouldBe200IfNotHeaderWritten(t *testing.T) {
	resp := newWrappedWriter(httptest.NewRecorder())
	req := httptest.NewRequest("GET", "/api/v4/test", nil)
	handler := TestHandler{func(w http.ResponseWriter, r *http.Request) {
		n, err := w.Write([]byte{})
		require.NoError(t, err, "Failed to write response")
		require.Equal(t, 0, n, "Expected to write 0 bytes")
	}}
	handler.ServeHTTP(resp, req)
	assert.Equal(t, http.StatusOK, resp.StatusCode())
}

func TestForUnsupportedHijack(t *testing.T) {
	resp := newWrappedWriter(httptest.NewRecorder())
	req := httptest.NewRequest("GET", "/api/v4/test", nil)
	handler := TestHandler{func(w http.ResponseWriter, r *http.Request) {
		conn, rw, err := w.(*responseWriterWrapper).Hijack()
		assert.Error(t, err)
		assert.Equal(t, "Hijacker interface not supported by the wrapped ResponseWriter", err.Error())
		assert.Nil(t, conn, "Expected nil connection")
		assert.Nil(t, rw, "Expected nil buffer")
	}}
	handler.ServeHTTP(resp, req)
}

func TestForSupportedHijack(t *testing.T) {
	resp := newWrappedWriter(newResponseWithHijack(httptest.NewRecorder()))
	req := httptest.NewRequest("GET", "/api/v4/test", nil)
	handler := TestHandler{func(w http.ResponseWriter, r *http.Request) {
		conn, rw, err := w.(*responseWriterWrapper).Hijack()
		require.NoError(t, err, "Hijack should succeed with supporting ResponseWriter")
		assert.Nil(t, conn, "Expected nil connection from test implementation")
		assert.Nil(t, rw, "Expected nil buffer from test implementation")
	}}
	handler.ServeHTTP(resp, req)
}

func TestForSupportedFlush(t *testing.T) {
	resp := newWrappedWriter(httptest.NewRecorder())
	req := httptest.NewRequest("GET", "/api/v4/test", nil)
	handler := TestHandler{func(w http.ResponseWriter, r *http.Request) {
		n, err := w.Write([]byte{})
		require.NoError(t, err, "Failed to write response")
		require.Equal(t, 0, n, "Expected to write 0 bytes")
		w.(*responseWriterWrapper).Flush()
	}}
	handler.ServeHTTP(resp, req)
	assert.Equal(t, http.StatusOK, resp.StatusCode())
}