70 lines
1.6 KiB
Go
70 lines
1.6 KiB
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 }
|