Test register endpoint

internal/builds/register.go

@@ -12,7 +12,7 @@ import (
 )
 
 const (
-	maxRegisterBodySize = 4 * 1024
+	maxRegisterBodySize = 32 * 1024
 	runnerBuildQueue    = "runner:build_queue:"
 )
 
@@ -36,6 +36,8 @@ var (
 type largeBodyError struct{ error }
 type watchError struct{ error }
 
+type WatchKeyHandler func(key, value string, timeout time.Duration) (redis.WatchKeyStatus, error)
+
 func init() {
 	prometheus.MustRegister(
 		registerHandlerHits,
@@ -77,14 +79,14 @@ func proxyRegisterRequest(h http.Handler, w http.ResponseWriter, r *http.Request
 	h.ServeHTTP(w, r)
 }
 
-func watchForRunnerChange(token, lastUpdate string, duration time.Duration) (redis.WatchKeyStatus, error) {
+func watchForRunnerChange(watchHandler WatchKeyHandler, token, lastUpdate string, duration time.Duration) (redis.WatchKeyStatus, error) {
 	registerHandlerOpen.WithLabelValues("watching").Inc()
 	defer registerHandlerOpen.WithLabelValues("watching").Dec()
 
-	return redis.WatchKey(runnerBuildQueue+token, lastUpdate, duration)
+	return watchHandler(runnerBuildQueue+token, lastUpdate, duration)
 }
 
-func RegisterHandler(h http.Handler, pollingDuration time.Duration) http.Handler {
+func RegisterHandler(h http.Handler, watchHandler WatchKeyHandler, pollingDuration time.Duration) http.Handler {
 	if pollingDuration == 0 {
 		return h
 	}
@@ -97,23 +99,26 @@ func RegisterHandler(h http.Handler, pollingDuration time.Duration) http.Handler
 			return
 		}
 
+		newRequest := helper.CloneRequestWithNewBody(r, requestBody)
+
 		runnerRequest, err := readRunnerRequest(r, requestBody)
 		if err != nil {
 			registerHandlerHits.WithLabelValues("body-parse-error").Inc()
-			proxyRegisterRequest(h, w, r)
+			proxyRegisterRequest(h, w, newRequest)
 			return
 		}
 
 		if runnerRequest.Token == "" || runnerRequest.LastUpdate == "" {
 			registerHandlerHits.WithLabelValues("missing-values").Inc()
-			proxyRegisterRequest(h, w, r)
+			proxyRegisterRequest(h, w, newRequest)
 			return
 		}
 
-		result, err := watchForRunnerChange(runnerRequest.Token, runnerRequest.LastUpdate, pollingDuration)
+		result, err := watchForRunnerChange(watchHandler, runnerRequest.Token,
+			runnerRequest.LastUpdate, pollingDuration)
 		if err != nil {
 			registerHandlerHits.WithLabelValues("watch-error").Inc()
-			helper.Fail500(w, r, &watchError{err})
+			helper.Fail500(w, newRequest, &watchError{err})
 			return
 		}
 
@@ -122,7 +127,7 @@ func RegisterHandler(h http.Handler, pollingDuration time.Duration) http.Handler
 		// We proxy request to Rails, to see whether we can receive the build
 		case redis.WatchKeyStatusAlreadyChanged:
 			registerHandlerHits.WithLabelValues("already-changed").Inc()
-			proxyRegisterRequest(h, w, helper.CloneRequestWithNewBody(r, requestBody))
+			proxyRegisterRequest(h, w, newRequest)
 
 		// It means that we detected a change after watching.
 		// We could potentially proxy request to Rails, but...

internal/builds/register_test.go

+package builds
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"errors"
+	"github.com/stretchr/testify/assert"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/redis"
+)
+
+func echoRequest(rw http.ResponseWriter, req *http.Request) {
+	io.Copy(rw, req.Body)
+}
+
+var echoRequestFunc = http.HandlerFunc(echoRequest)
+
+func TestRegisterHandlerLargeBody(t *testing.T) {
+	h := RegisterHandler(echoRequestFunc, nil, time.Second)
+
+	data := make([]byte, maxRegisterBodySize+5)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("POST", "/", bytes.NewBuffer(data))
+
+	h.ServeHTTP(rw, req)
+
+	assert.Equal(t, http.StatusInternalServerError, rw.Code)
+}
+
+func TestRegisterHandlerInvalidRunnerRequest(t *testing.T) {
+	h := RegisterHandler(echoRequestFunc, nil, time.Second)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("POST", "/", bytes.NewBufferString("invalid"))
+
+	h.ServeHTTP(rw, req)
+
+	assert.Equal(t, http.StatusOK, rw.Code)
+	assert.Equal(t, "invalid", rw.Body.String())
+}
+
+func TestRegisterHandlerInvalidJsonPayload(t *testing.T) {
+	h := RegisterHandler(echoRequestFunc, nil, time.Second)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("POST", "/", bytes.NewBufferString("{["))
+	req.Header.Set("Content-Type", "application/json")
+
+	h.ServeHTTP(rw, req)
+
+	assert.Equal(t, http.StatusOK, rw.Code)
+	assert.Equal(t, "{[", rw.Body.String())
+}
+
+func TestRegisterHandlerMissingData(t *testing.T) {
+	datas := []string{"{\"token\":\"token\"}", "{\"last_update\":\"data\"}"}
+
+	for _, data := range datas {
+		h := RegisterHandler(echoRequestFunc, nil, time.Second)
+
+		rw := httptest.NewRecorder()
+		req := httptest.NewRequest("POST", "/", bytes.NewBufferString(data))
+		req.Header.Set("Content-Type", "application/json")
+
+		h.ServeHTTP(rw, req)
+
+		assert.Equal(t, http.StatusOK, rw.Code)
+		assert.Equal(t, data, rw.Body.String())
+	}
+}
+
+func TestRegisterHandlerWatcherError(t *testing.T) {
+	data := "{\"token\":\"token\",\"last_update\":\"last_update\"}"
+
+	executed := false
+	watchKeyHandler := func(key, value string, timeout time.Duration) (redis.WatchKeyStatus, error) {
+		executed = true
+		return redis.WatchKeyStatusNoChange, errors.New("redis connection")
+	}
+
+	h := RegisterHandler(echoRequestFunc, watchKeyHandler, time.Second)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("POST", "/", bytes.NewBufferString(data))
+	req.Header.Set("Content-Type", "application/json")
+
+	h.ServeHTTP(rw, req)
+
+	assert.Equal(t, http.StatusInternalServerError, rw.Code)
+	assert.True(t, executed)
+}
+
+func TestRegisterHandlerWatcherAlreadyChanged(t *testing.T) {
+	data := "{\"token\":\"token\",\"last_update\":\"last_update\"}"
+
+	executed := false
+	watchKeyHandler := func(key, value string, timeout time.Duration) (redis.WatchKeyStatus, error) {
+		assert.Equal(t, "runner:build_queue:token", key)
+		assert.Equal(t, "last_update", value)
+		executed = true
+		return redis.WatchKeyStatusAlreadyChanged, nil
+	}
+
+	h := RegisterHandler(echoRequestFunc, watchKeyHandler, time.Second)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("POST", "/", bytes.NewBufferString(data))
+	req.Header.Set("Content-Type", "application/json")
+
+	h.ServeHTTP(rw, req)
+
+	assert.Equal(t, http.StatusOK, rw.Code)
+	assert.Equal(t, data, rw.Body.String())
+	assert.True(t, executed)
+}
+
+func TestRegisterHandlerWatcherSeenChange(t *testing.T) {
+	data := "{\"token\":\"token\",\"last_update\":\"last_update\"}"
+
+	executed := false
+	watchKeyHandler := func(key, value string, timeout time.Duration) (redis.WatchKeyStatus, error) {
+		assert.Equal(t, "runner:build_queue:token", key)
+		assert.Equal(t, "last_update", value)
+		executed = true
+		return redis.WatchKeyStatusSeenChange, nil
+	}
+
+	h := RegisterHandler(echoRequestFunc, watchKeyHandler, time.Second)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("POST", "/", bytes.NewBufferString(data))
+	req.Header.Set("Content-Type", "application/json")
+
+	h.ServeHTTP(rw, req)
+
+	assert.Equal(t, http.StatusNoContent, rw.Code)
+	assert.True(t, executed)
+}
+
+func TestRegisterHandlerWatcherTimeout(t *testing.T) {
+	data := "{\"token\":\"token\",\"last_update\":\"last_update\"}"
+
+	executed := false
+	watchKeyHandler := func(key, value string, timeout time.Duration) (redis.WatchKeyStatus, error) {
+		assert.Equal(t, "runner:build_queue:token", key)
+		assert.Equal(t, "last_update", value)
+		executed = true
+		return redis.WatchKeyStatusTimeout, nil
+	}
+
+	h := RegisterHandler(echoRequestFunc, watchKeyHandler, time.Second)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("POST", "/", bytes.NewBufferString(data))
+	req.Header.Set("Content-Type", "application/json")
+
+	h.ServeHTTP(rw, req)
+
+	assert.Equal(t, http.StatusNoContent, rw.Code)
+	assert.True(t, executed)
+}
+
+func TestRegisterHandlerWatcherNoChange(t *testing.T) {
+	data := "{\"token\":\"token\",\"last_update\":\"last_update\"}"
+
+	executed := false
+	watchKeyHandler := func(key, value string, timeout time.Duration) (redis.WatchKeyStatus, error) {
+		assert.Equal(t, "runner:build_queue:token", key)
+		assert.Equal(t, "last_update", value)
+		executed = true
+		return redis.WatchKeyStatusNoChange, nil
+	}
+
+	h := RegisterHandler(echoRequestFunc, watchKeyHandler, time.Second)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("POST", "/", bytes.NewBufferString(data))
+	req.Header.Set("Content-Type", "application/json")
+
+	h.ServeHTTP(rw, req)
+
+	assert.Equal(t, http.StatusNoContent, rw.Code)
+	assert.True(t, executed)
+}

internal/redis/redis.go

@@ -135,6 +135,10 @@ func Configure(cfg *config.RedisConfig, dialFunc func() (redis.Conn, error)) {
 	}
 }
 
+func Unconfigure() {
+	pool = nil
+}
+
 // Get a connection for the Redis-pool
 func Get() redis.Conn {
 	if pool != nil {