feat: add Responder.Then() method

Signed-off-by: Maxime Soulé <btik-git@scoubidou.com>

Author: maxatome

response.go

@@ -2,11 +2,14 @@ package httpmock
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"encoding/xml"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
+	"reflect"
 	"strconv"
 	"strings"
 	"sync"
@@ -122,6 +125,79 @@ func (r Responder) Delay(d time.Duration) Responder {
 	}
 }
 
+type fromThenKeyType struct{}
+
+var fromThenKey = fromThenKeyType{}
+
+var errThenDone = errors.New("ThenDone")
+
+// similar is simple but a bit tricky. Here we consider two Responder
+// are similar if they share the same function, but not necessarily
+// the same environment. It is only used by Then below.
+func (r Responder) similar(other Responder) bool {
+	return reflect.ValueOf(r).Pointer() == reflect.ValueOf(other).Pointer()
+}
+
+// Then returns a new Responder that calls r on first invocation, then
+// next on following ones, except when Then is chained, in this case
+// next is called only once:
+//   A := httpmock.NewStringResponder(200, "A")
+//   B := httpmock.NewStringResponder(200, "B")
+//   C := httpmock.NewStringResponder(200, "C")
+//
+//   httpmock.RegisterResponder("GET", "/pipo", A.Then(B).Then(C))
+//
+//   http.Get("http://foo.bar/pipo") // A is called
+//   http.Get("http://foo.bar/pipo") // B is called
+//   http.Get("http://foo.bar/pipo") // C is called
+//   http.Get("http://foo.bar/pipo") // C is called, and so on
+//
+// A panic occurs if next is the result of another Then call (because
+// allowing it could cause inextricable problems at runtime). Then
+// calls can be chained, but cannot call each other by
+// parameter. Example:
+//   A.Then(B).Then(C) // is OK
+//   A.Then(B.Then(C)) // panics as A.Then() parameter is another Then() call
+func (r Responder) Then(next Responder) (x Responder) {
+	var done int
+	var mu sync.Mutex
+	x = func(req *http.Request) (*http.Response, error) {
+		mu.Lock()
+		defer mu.Unlock()
+
+		ctx := req.Context()
+		thenCalledUs, _ := ctx.Value(fromThenKey).(bool)
+		if !thenCalledUs {
+			req = req.WithContext(context.WithValue(ctx, fromThenKey, true))
+		}
+
+		switch done {
+		case 0:
+			resp, err := r(req)
+			if err != errThenDone {
+				if !x.similar(r) { // r is NOT a Then
+					done = 1
+				}
+				return resp, err
+			}
+			fallthrough
+
+		case 1:
+			done = 2 // next is NEVER a Then, as it is forbidden by design
+			return next(req)
+		}
+		if thenCalledUs {
+			return nil, errThenDone
+		}
+		return next(req)
+	}
+
+	if next.similar(x) {
+		panic("Then() does not accept another Then() Responder as parameter")
+	}
+	return
+}
+
 // ResponderFromResponse wraps an *http.Response in a Responder.
 //
 // Be careful, except for responses generated by httpmock

response_test.go

@@ -655,6 +655,98 @@ func TestResponder(t *testing.T) {
 	}
 }
 
+func TestResponder_Then(t *testing.T) {
+	req, err := http.NewRequest(http.MethodGet, "http://foo.bar", nil)
+	if err != nil {
+		t.Fatalf("Error creating request: %s", err)
+	}
+
+	//
+	// Then
+	var stack string
+	newResponder := func(level string) Responder {
+		return func(*http.Request) (*http.Response, error) {
+			stack += level
+			return NewStringResponse(200, level), nil
+		}
+	}
+	var rt Responder
+	chk := func(t *testing.T, expectedLevel, expectedStack string) {
+		helper(t).Helper()
+		resp, err := rt(req)
+		if err != nil {
+			t.Errorf("Responder retruned an unexpected error: %s", err)
+			return
+		}
+		b, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			t.Errorf("Read response failed: %s", err)
+			return
+		}
+		if string(b) != expectedLevel {
+			t.Errorf("level: got %q but expected %q (stack is %q)",
+				b, expectedLevel, stack)
+			return
+		}
+		if stack != expectedStack {
+			t.Errorf("stack: got %q but expected %q", stack, expectedStack)
+			return
+		}
+	}
+
+	A, B, C := newResponder("A"), newResponder("B"), newResponder("C")
+	D, E, F := newResponder("D"), newResponder("E"), newResponder("F")
+
+	t.Run("simple", func(t *testing.T) {
+		// (r=A,then=B)
+		rt = A.Then(B)
+
+		chk(t, "A", "A")
+		chk(t, "B", "AB")
+		chk(t, "B", "ABB")
+		chk(t, "B", "ABBB")
+	})
+
+	stack = ""
+
+	t.Run("simple chained", func(t *testing.T) {
+		//             (r=A,then=B)
+		//          (r=↑,then=C)
+		//       (r=↑,then=D)
+		//    (r=↑,then=E)
+		// (r=↑,then=F)
+		rt = A.Then(B).
+			Then(C).
+			Then(D).
+			Then(E).
+			Then(F)
+
+		chk(t, "A", "A")
+		chk(t, "B", "AB")
+		chk(t, "C", "ABC")
+		chk(t, "D", "ABCD")
+		chk(t, "E", "ABCDE")
+		chk(t, "F", "ABCDEF")
+		chk(t, "F", "ABCDEFF")
+		chk(t, "F", "ABCDEFFF")
+	})
+
+	stack = ""
+
+	t.Run("Then Responder as Then param", func(t *testing.T) {
+		panicked, str := catchPanic(func() {
+			A.Then(B.Then(C))
+		})
+		if str != "Then() does not accept another Then() Responder as parameter" {
+			if !panicked {
+				t.Error("Should have panicked")
+			} else {
+				t.Errorf("Wrong panic message: %q", str)
+			}
+		}
+	})
+}
+
 func TestParallelResponder(t *testing.T) {
 	req, err := http.NewRequest(http.MethodGet, "http://foo.bar", nil)
 	if err != nil {