add until functionality (#11)

Co-authored-by: Dean Oren <[email protected]>

pkg/retry/retry.go

@@ -9,10 +9,11 @@ import (
 )
 
 type Retry struct {
-	Timeout        time.Duration          // max duration
-	MaxRetries     *int                   // max retries, when nil, there's no retry limit
-	Throttle       time.Duration          // wait duration between calls
-	IsRetryableFns []func(err error) bool // functions to determine if error is retriable
+	Timeout        time.Duration               // max duration
+	MaxRetries     *int                        // max retries, when nil, there's no retry limit
+	Throttle       time.Duration               // wait duration between calls
+	IsRetryableFns []func(err error) bool      // functions to determine if error is retriable
+	UntilFns       []func(*http.Response) bool // functions to the determine if the given response is the expected response
 }
 
 type Config interface {
@@ -42,6 +43,8 @@ func (c *Retry) withConfig(cfgs ...Config) *Retry {
 			c.MaxRetries = cfg.Value().(*int)
 		case CONFIG_IS_RETRYABLE:
 			c.IsRetryableFns = cfg.Value().([]func(err error) bool)
+		case CONFIG_UNTIL:
+			c.UntilFns = cfg.Value().([]func(*http.Response) bool)
 		}
 	}
 	return c
@@ -67,7 +70,11 @@ func (r *Retry) Do(req *http.Request, do func(*http.Request) (*http.Response, er
 
 	for {
 		res, maxRetries, retry, lastErr = r.doIteration(newReq, do, maxRetries)
-		if lastErr == nil || !retry {
+		if lastErr == nil && r.isFulfilled(res) {
+			return res, nil
+		}
+
+		if !retry {
 			return res, lastErr
 		}
 
@@ -84,13 +91,13 @@ func (r *Retry) Do(req *http.Request, do func(*http.Request) (*http.Response, er
 	}
 }
 
+// doIteration runs the do function with the given request
 func (r *Retry) doIteration(req *http.Request, do func(*http.Request) (*http.Response, error), retries int) (res *http.Response, retriesLeft int, retry bool, err error) {
 	retriesLeft = retries
 	retry = true
 
 	res, err = do(req)
 	if err == nil {
-		retry = false
 		return
 	}
 
@@ -113,3 +120,13 @@ func (r *Retry) doIteration(req *http.Request, do func(*http.Request) (*http.Res
 
 	return
 }
+
+// isFulfilled check if Until functions are all fulfilled
+func (r *Retry) isFulfilled(res *http.Response) bool {
+	for _, f := range r.UntilFns {
+		if !f(res) {
+			return false
+		}
+	}
+	return true
+}

pkg/retry/retry_test.go

@@ -136,3 +136,42 @@ func TestClient_DoWithRetryTimeout(t *testing.T) {
 		t.Errorf("expected do request to return retry context timed out error but got '%v' instead", err)
 	}
 }
+
+func TestClient_DoWithUntil(t *testing.T) {
+	c, mux, teardown, err := client.MockServer()
+	defer teardown()
+	if err != nil {
+		t.Errorf("error from mock.AuthServer: %s", err.Error())
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
+	defer cancel()
+	mux.HandleFunc("/2s", func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		if ctx.Err() != nil {
+			w.WriteHeader(http.StatusOK)
+			fmt.Fprint(w, response_after_2s)
+		} else {
+			w.WriteHeader(http.StatusOK)
+			fmt.Fprint(w, response_before_2s)
+		}
+	})
+
+	var got struct {
+		Status bool `json:"status"`
+	}
+
+	c = c.WithRetry(retry.New().SetThrottle(1 * time.Second).Until(func(r *http.Response) bool {
+		return got.Status
+	}))
+
+	req, _ := c.Request(context.Background(), http.MethodGet, "/2s", nil)
+
+	if _, err := c.Do(req, &got); err != nil {
+		t.Errorf("do request: %v", err)
+	}
+
+	if !got.Status {
+		t.Errorf("received status = %v", got.Status)
+	}
+}

pkg/retry/unti_test.go

@@ -0,0 +1,42 @@
+package retry
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestRetry_Until(t *testing.T) {
+	r := New()
+	r.UntilFns = []func(*http.Response) bool{noOp}
+
+	type args struct {
+		f []func(*http.Response) bool
+	}
+	tests := []struct {
+		name string
+		args args
+		want *Retry
+	}{
+		{"ok", args{[]func(*http.Response) bool{noOp}}, r},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := New()
+			got := c.Until(tt.args.f...)
+			if len(got.UntilFns) != len(r.UntilFns) {
+				t.Error("wrong lengths")
+				return
+			}
+			for k, v := range got.UntilFns {
+				if GetFunctionName(r.UntilFns[k]) != GetFunctionName(v) {
+					t.Errorf("%s != %s", GetFunctionName(r.UntilFns[k]), GetFunctionName(v))
+					return
+				}
+			}
+		})
+	}
+}
+
+func noOp(_ *http.Response) bool {
+	return true
+}

pkg/retry/until.go

@@ -0,0 +1,31 @@
+package retry
+
+import "net/http"
+
+const (
+	// Until configuration constants
+	CONFIG_UNTIL = "UntilFns"
+)
+
+// UntilFns is the config struct
+type UntilFns struct {
+	fnList []func(*http.Response) bool
+}
+
+// Until sets functions that determin if the response given is the expected response
+// this functionality is useful, for example, when waiting for an API status to be ready
+func (c *Retry) Until(f ...func(*http.Response) bool) *Retry {
+	return c.withConfig(&UntilFns{
+		fnList: f,
+	})
+}
+
+var _ = Config(&UntilFns{})
+
+func (c *UntilFns) String() string {
+	return CONFIG_UNTIL
+}
+
+func (c *UntilFns) Value() interface{} {
+	return c.fnList
+}