From d323c2089660a20b1c240ceb1c6042a252c1b985 Mon Sep 17 00:00:00 2001
From: Renan DelValle <renanidelvalle+noreply@gmail.com>
Date: Tue, 30 Jan 2018 17:45:36 -0800
Subject: [PATCH] Fixing logic that can lead to nil error being returned and
 retry stopping early.

---
 realis_e2e_test.go | 26 ++++++++++++++++++++++++++
 retry.go           | 15 ++++++++++++---
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/realis_e2e_test.go b/realis_e2e_test.go
index 1538886..d5c2b4e 100644
--- a/realis_e2e_test.go
+++ b/realis_e2e_test.go
@@ -58,6 +58,32 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
+func TestBadEndpoint(t *testing.T) {
+
+	// Attempt to connect to a bad endpoint
+	r, err := realis.NewRealisClient(realis.SchedulerUrl("http://192.168.33.7:8081/scheduler/"),
+		realis.TimeoutMS(200),
+		realis.BackOff(&realis.Backoff{ // Reduce penalties for this test to make it quick
+			Steps:    5,
+			Duration: 1 * time.Second,
+			Factor:   1.0,
+			Jitter:   0.1}),
+	)
+	defer r.Close()
+
+	taskQ := &aurora.TaskQuery{
+		Role:        "no",
+		Environment: "task",
+		JobName:     "here",
+	}
+
+	_, err = r.GetTasksWithoutConfigs(taskQ)
+
+	// Check that we do error out of retrying
+	assert.Error(t, err)
+
+}
+
 func TestLeaderFromZK(t *testing.T) {
 	cluster := realis.GetDefaultClusterFromZKUrl("192.168.33.7:2181")
 	url, err := realis.LeaderFromZK(*cluster)
diff --git a/retry.go b/retry.go
index dd7b113..073242b 100644
--- a/retry.go
+++ b/retry.go
@@ -17,10 +17,11 @@ limitations under the License.
 package realis
 
 import (
-	"errors"
 	"time"
 
 	"math/rand"
+
+	"github.com/pkg/errors"
 )
 
 // Jitter returns a time.Duration between duration and duration + maxFactor *
@@ -52,6 +53,8 @@ type ConditionFunc func() (done bool, err error)
 // If the condition never returns true, ErrWaitTimeout is returned. All other
 // errors terminate immediately.
 func ExponentialBackoff(backoff Backoff, condition ConditionFunc) error {
+	var err error
+	var ok bool
 	duration := backoff.Duration
 	for i := 0; i < backoff.Steps; i++ {
 		if i != 0 {
@@ -63,7 +66,7 @@ func ExponentialBackoff(backoff Backoff, condition ConditionFunc) error {
 			duration = time.Duration(float64(duration) * backoff.Factor)
 		}
 
-		ok, err := condition()
+		ok, err = condition()
 
 		// If the function executed says it succeeded, stop retrying
 		if ok {
@@ -78,5 +81,11 @@ func ExponentialBackoff(backoff Backoff, condition ConditionFunc) error {
 		}
 
 	}
-	return NewTimeoutError(errors.New("Timed out while retrying"))
+
+	// Provide more information to the user wherever possible
+	if err != nil {
+		return NewTimeoutError(errors.Wrap(err, "Timed out while retrying"))
+	} else {
+		return NewTimeoutError(errors.New("Timed out while retrying"))
+	}
 }