gitlab-org--gitlab-foss/workhorse/internal/senddata/writer_test.go

package senddata

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

	"gitlab.com/gitlab-org/gitlab/workhorse/internal/headers"
)

func TestWriter(t *testing.T) {
	upstreamResponse := "hello world"

	testCases := []struct {
		desc        string
		headerValue string
		out         string
	}{
		{
			desc:        "inject",
			headerValue: testInjecterName + ":" + testInjecterName,
			out:         testInjecterData,
		},
		{
			desc:        "pass",
			headerValue: "",
			out:         upstreamResponse,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.desc, func(t *testing.T) {
			recorder := httptest.NewRecorder()
			rw := &sendDataResponseWriter{rw: recorder, injecters: []Injecter{&testInjecter{}}}

			rw.Header().Set(headers.GitlabWorkhorseSendDataHeader, tc.headerValue)

			n, err := rw.Write([]byte(upstreamResponse))
			require.NoError(t, err)
			require.Equal(t, len(upstreamResponse), n, "bytes written")

			recorder.Flush()

			body := recorder.Result().Body
			data, err := io.ReadAll(body)
			require.NoError(t, err)
			require.NoError(t, body.Close())

			require.Equal(t, tc.out, string(data))
		})
	}
}

const (
	testInjecterName = "test-injecter"
	testInjecterData = "hello this is injected data"
)

type testInjecter struct{}

func (ti *testInjecter) Inject(w http.ResponseWriter, r *http.Request, sendData string) {
	io.WriteString(w, testInjecterData)
}

func (ti *testInjecter) Match(s string) bool { return strings.HasPrefix(s, testInjecterName+":") }
func (ti *testInjecter) Name() string        { return testInjecterName }