[Bugfix] Retry on 503 in tests (#311)

Author: ajanikow

.travis.yml

@@ -28,7 +28,7 @@ services:
 language: go
 
 env:
-  - TEST_SUITE=run-unit-tests ALWAYS=1
+  - TEST_SUITE=run-unit-tests GOIMAGE=gcr.io/gcr-for-testing/golang:1.13.4-stretch ALWAYS=1
   - TEST_SUITE=run-tests-single GOIMAGE=gcr.io/gcr-for-testing/golang:1.13.4-stretch STARTER=gcr.io/gcr-for-testing/arangodb/arangodb-starter:latest ALPINE_IMAGE=gcr.io/gcr-for-testing/alpine:3.4 ARANGODB=gcr.io/gcr-for-testing/arangodb:3.6
   - TEST_SUITE=run-tests-single GOIMAGE=gcr.io/gcr-for-testing/golang:1.13.4-stretch STARTER=gcr.io/gcr-for-testing/arangodb/arangodb-starter:latest ALPINE_IMAGE=gcr.io/gcr-for-testing/alpine:3.4 ARANGODB=gcr.io/gcr-for-testing/arangodb/arangodb:latest ALWAYS=1
   - TEST_SUITE=run-tests-single GOIMAGE=gcr.io/gcr-for-testing/golang:1.13.4-stretch STARTER=gcr.io/gcr-for-testing/arangodb/arangodb-starter:latest ALPINE_IMAGE=gcr.io/gcr-for-testing/alpine:3.4 ARANGODB=gcr.io/gcr-for-testing/arangodb:3.7

error.go

@@ -100,6 +100,16 @@ func IsArangoError(err error) bool {
 	return ok && ae.HasError
 }
 
+// AsArangoError returns true when the given error is an ArangoError together with an object.
+func AsArangoError(err error) (ArangoError, bool) {
+	ae, ok := Cause(err).(ArangoError)
+	if ok {
+		return ae, true
+	} else {
+		return ArangoError{}, false
+	}
+}
+
 // IsArangoErrorWithCode returns true when the given error is an ArangoError and its Code field is equal to the given code.
 func IsArangoErrorWithCode(err error, code int) bool {
 	ae, ok := Cause(err).(ArangoError)

test/agency_test.go

@@ -42,6 +42,19 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func checkAgencyEndpoints(ctx context.Context, c driver.Client) error {
+	cl, err := c.Cluster(ctx)
+	if err != nil {
+		return err
+	}
+	_, err = cl.Health(ctx)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // getAgencyEndpoints queries the cluster to get all agency endpoints.
 func getAgencyEndpoints(ctx context.Context, c driver.Client) ([]string, error) {
 	cl, err := c.Cluster(ctx)
@@ -114,6 +127,11 @@ func getAgencyConnection(ctx context.Context, t testEnv, c driver.Client) (agenc
 		// These tests assume an HTTP connetion, so we skip under this condition
 		return nil, driver.ArangoError{HasError: true, Code: 412, ErrorMessage: "Using vst is not supported in agency tests"}
 	}
+
+	if err := driverErrorCheck(ctx, c, checkAgencyEndpoints, driverErrorCheckRetry503).Retry(125*time.Millisecond, time.Minute); err != nil {
+		return nil, err
+	}
+
 	endpoints, err := getAgencyEndpoints(ctx, c)
 	if err != nil {
 		return nil, err

test/check_test.go

@@ -0,0 +1,69 @@
+//
+// DISCLAIMER
+//
+// Copyright 2021 ArangoDB GmbH, Cologne, Germany
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Copyright holder is ArangoDB GmbH, Cologne, Germany
+//
+// Author Adam Janikowski
+//
+
+package test
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/arangodb/go-driver"
+)
+
+type driverErrorCheckFunc func(err error) (bool, error)
+type driverErrorChecker func(ctx context.Context, client driver.Client) error
+
+func driverErrorCheck(ctx context.Context, c driver.Client, checker driverErrorChecker, checks ...driverErrorCheckFunc) retryFunc {
+	return func() error {
+		err := checker(ctx, c)
+
+		for _, check := range checks {
+			if valid, err := check(err); err != nil {
+				return err
+			} else if !valid {
+				return nil
+			}
+		}
+
+		return interrupt{}
+	}
+}
+
+func driverErrorCheckRetry503(err error) (bool, error) {
+	if err == nil {
+		return true, nil
+	}
+
+	if ae, ok := driver.AsArangoError(err); !ok {
+		return false, err
+	} else {
+		if !ae.HasError {
+			return true, nil
+		}
+		switch ae.Code {
+		case http.StatusServiceUnavailable:
+			return false, nil
+		default:
+			return true, err
+		}
+	}
+}