Merge branch 'cat-geo-middleware-for-workhorse-tests' into 'master'

Introduce Geo proxy bypass handler for Workhorse tests

See merge request gitlab-org/gitlab!74908

workhorse/internal/api/api_test.go

@@ -49,7 +49,7 @@ func getGeoProxyURLGivenResponse(t *testing.T, givenInternalApiResponse string)
 }
 
 func testRailsServer(url *regexp.Regexp, code int, body string) *httptest.Server {
-	return testhelper.TestServerWithHandler(url, func(w http.ResponseWriter, r *http.Request) {
+	return testhelper.TestServerWithHandlerWithGeoPolling(url, func(w http.ResponseWriter, r *http.Request) {
 		// return a 204 No Content response if we don't receive the JWT header
 		if r.Header.Get(secret.RequestHeader) == "" {
 			w.WriteHeader(204)

workhorse/internal/testhelper/testhelper.go

@@ -22,6 +22,10 @@ import (
 	"gitlab.com/gitlab-org/gitlab/workhorse/internal/secret"
 )
 
+const (
+	geoProxyEndpointPath = "/api/v4/geo/proxy"
+)
+
 func ConfigureSecret() {
 	secret.SetPath(path.Join(RootDir(), "testdata/test-secret"))
 }
@@ -50,7 +54,20 @@ func RequireResponseHeader(t *testing.T, w interface{}, header string, expected
 	require.Equal(t, expected, actual, "values for HTTP header %s", header)
 }
 
+// TestServerWithHandler skips Geo API polling for a proxy URL by default,
+// use TestServerWithHandlerWithGeoPolling if you need to explicitly
+// handle Geo API polling request as well.
 func TestServerWithHandler(url *regexp.Regexp, handler http.HandlerFunc) *httptest.Server {
+	return TestServerWithHandlerWithGeoPolling(url, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == geoProxyEndpointPath {
+			return
+		}
+
+		handler(w, r)
+	}))
+}
+
+func TestServerWithHandlerWithGeoPolling(url *regexp.Regexp, handler http.HandlerFunc) *httptest.Server {
 	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		logEntry := log.WithFields(log.Fields{
 			"method": r.Method,

workhorse/internal/upstream/upstream.go

@@ -207,8 +207,8 @@ func (u *upstream) findGeoProxyRoute(cleanedPath string, r *http.Request) *route
 
 func (u *upstream) pollGeoProxyAPI() {
 	for {
-		u.geoProxyPollSleep(geoProxyApiPollingInterval)
 		u.callGeoProxyAPI()
+		u.geoProxyPollSleep(geoProxyApiPollingInterval)
 	}
 }
 

workhorse/internal/upstream/upstream_test.go

@@ -314,9 +314,5 @@ func startWorkhorseServer(railsServerURL string, enableGeoProxyFeature bool) (*h
 		}
 	}
 
-	// Since the first sleep happens before any API call, this ensures
-	// we call the API at least once.
-	waitForNextApiPoll()
-
 	return ws, ws.Close, waitForNextApiPoll
 }

workhorse/sendfile_test.go

@@ -5,13 +5,15 @@ import (
 	"io/ioutil"
 	"mime"
 	"net/http"
-	"net/http/httptest"
 	"os"
 	"path"
+	"regexp"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/labkit/log"
+
+	"gitlab.com/gitlab-org/gitlab/workhorse/internal/testhelper"
 )
 
 func TestDeniedLfsDownload(t *testing.T) {
@@ -35,7 +37,7 @@ func allowedXSendfileDownload(t *testing.T, contentFilename string, filePath str
 	prepareDownloadDir(t)
 
 	// Prepare test server and backend
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	ts := testhelper.TestServerWithHandler(regexp.MustCompile(`.`), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		log.WithFields(log.Fields{"method": r.Method, "url": r.URL}).Info("UPSTREAM")
 
 		require.Equal(t, "X-Sendfile", r.Header.Get("X-Sendfile-Type"))
@@ -69,7 +71,7 @@ func deniedXSendfileDownload(t *testing.T, contentFilename string, filePath stri
 	prepareDownloadDir(t)
 
 	// Prepare test server and backend
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	ts := testhelper.TestServerWithHandler(regexp.MustCompile(`.`), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		log.WithFields(log.Fields{"method": r.Method, "url": r.URL}).Info("UPSTREAM")
 
 		require.Equal(t, "X-Sendfile", r.Header.Get("X-Sendfile-Type"))