Add Status to GitLab client

2 files changed, 99 insertions, 1 deletions

diff --git a/internal/source/gitlab/client/client.go b/internal/source/gitlab/client/client.go
index 3a805ce9..51e7eca0 100644
--- a/internal/source/gitlab/client/client.go
+++ b/internal/source/gitlab/client/client.go
@@ -18,6 +18,11 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/metrics"
 )
 
+// ConnectionErrorMsg to be returned with `gc.Status` if Pages
+// fails to connect to the internal GitLab API, times out
+// or a 401 given that the credentials used are wrong
+const ConnectionErrorMsg = "failed to connect to internal Pages API"
+
 // Client is a HTTP client to access Pages internal API
 type Client struct {
 	secretKey      []byte
@@ -100,6 +105,19 @@ func (gc *Client) GetLookup(ctx context.Context, host string) api.Lookup {
 	return lookup
 }
 
+// Status checks that Pages can reach the rails internal Pages API
+// for source domain configuration.
+// Timeout is the same as -gitlab-client-http-timeout
+func (gc *Client) Status() error {
+	resp, err := gc.get(context.Background(), "/api/v4/internal/pages/status", url.Values{})
+	if err != nil {
+		return fmt.Errorf("%s: %v", ConnectionErrorMsg, err)
+	}
+	defer resp.Body.Close()
+
+	return nil
+}
+
 func (gc *Client) get(ctx context.Context, path string, params url.Values) (*http.Response, error) {
 	endpoint, err := gc.endpoint(path, params)
 	if err != nil {
diff --git a/internal/source/gitlab/client/client_test.go b/internal/source/gitlab/client/client_test.go
index 02684860..ab90b474 100644
--- a/internal/source/gitlab/client/client_test.go
+++ b/internal/source/gitlab/client/client_test.go
@@ -10,7 +10,7 @@ import (
 	"testing"
 	"time"
 
-	jwt "github.com/dgrijalva/jwt-go"
+	"github.com/dgrijalva/jwt-go"
 	"github.com/stretchr/testify/require"
 
 	"gitlab.com/gitlab-org/gitlab-pages/internal/fixture"
@@ -231,6 +231,86 @@ func TestGetVirtualDomainAuthenticatedRequest(t *testing.T) {
 	require.Equal(t, "mygroup/myproject/public/", lookupPath.Source.Path)
 }
 
+func TestClientStatus(t *testing.T) {
+	tests := []struct {
+		name    string
+		status  int
+		wantErr bool
+	}{
+		{
+			name:   "api_enabled",
+			status: http.StatusNoContent,
+		},
+		{
+			name:    "api_unauthorized",
+			status:  http.StatusUnauthorized,
+			wantErr: true,
+		},
+		{
+			name:    "server_error",
+			status:  http.StatusInternalServerError,
+			wantErr: true,
+		},
+		{
+			name:    "gateway_timeout",
+			status:  http.StatusGatewayTimeout,
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mux := http.NewServeMux()
+			mux.HandleFunc("/api/v4/internal/pages/status", func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(tt.status)
+			})
+
+			server := httptest.NewServer(mux)
+			defer server.Close()
+
+			client := defaultClient(t, server.URL)
+
+			err := client.Status()
+			if tt.wantErr {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), ConnectionErrorMsg)
+				return
+			}
+
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestClientStatusClientTimeout(t *testing.T) {
+	timeout := 3 * time.Millisecond
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/api/v4/internal/pages/status", func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(timeout * 3)
+
+		w.WriteHeader(http.StatusOK)
+	})
+
+	server := httptest.NewServer(mux)
+	defer server.Close()
+
+	client := defaultClient(t, server.URL)
+	client.httpClient.Timeout = timeout
+
+	err := client.Status()
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Client.Timeout")
+}
+
+func TestClientStatusConnectionRefused(t *testing.T) {
+	client := defaultClient(t, "http://localhost:1234")
+
+	err := client.Status()
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "connection refused")
+}
+
 func validateToken(t *testing.T, tokenString string) {
 	t.Helper()
 	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {