diff --git a/ghttp/handlers.go b/ghttp/handlers.go index a80b27ddf..29fcf3ffe 100644 --- a/ghttp/handlers.go +++ b/ghttp/handlers.go @@ -10,6 +10,7 @@ import ( "net/url" "reflect" "strings" + "sync/atomic" "github.com/golang/protobuf/proto" "github.com/onsi/gomega" @@ -401,3 +402,53 @@ func RespondWithJSONEncodedPtr(statusCode *int, object interface{}, optionalHead func RespondWithProto(statusCode int, message proto.Message, optionalHeader ...http.Header) http.HandlerFunc { return NewGHTTPWithGomega(gomega.Default).RespondWithProto(statusCode, message, optionalHeader...) } + +var noOpHandler = func(_ http.ResponseWriter, _ *http.Request) { + // empty function for Nop +} + +//RespondWithMultiple +//This function combines a set of handlers into a handler such that each handler gets called in succession until +//the last handler. In this case when the last handler is reached it will keep responding with this handler. +//In the case of the RoundRobinWithMultiple combination handler the once the last handler is reached it will +//start at the beginning again. This is useful for testing retry behaviour of clients. +//I.E. A sequence of 500 500 500 200.... should be good enough for a retry to succeed. +func RespondWithMultiple(handlers ...http.HandlerFunc) http.HandlerFunc { + var responseNumber int32 = 0 + if len(handlers) > 0 { + return func(w http.ResponseWriter, req *http.Request) { + responseNum := atomic.LoadInt32(&responseNumber) + handlerNumber := min(responseNum, int32(len(handlers)-1)) + handlers[handlerNumber](w, req) + atomic.AddInt32(&responseNumber, 1) + } + } + return noOpHandler +} + +//RoundRobinWithMultiple +//This function combines a set of handlers into a single handler such that each handler gets called in succession until +//the last handler. Once the last handler is reached it will start at the beginning. +//In the case of the RespondWithMultiple combination handler the once the last handler is reached it will +//keep responding with the last handler. +// This is useful for stress testing retry behaviour over time. +//I.E. A sequence of 200 200 500 500 200 repeating can be called by many routines at the same time to +//imitate more complex behavior +func RoundRobinWithMultiple(handlers ...http.HandlerFunc) http.HandlerFunc { + var responseNumber int32 = 0 + if len(handlers) > 0 { + return func(w http.ResponseWriter, req *http.Request) { + handlerNumber := atomic.LoadInt32(&responseNumber) % int32(len(handlers)) + handlers[handlerNumber](w, req) + atomic.AddInt32(&responseNumber, 1) + } + } + return noOpHandler +} + +func min(one, two int32) int32 { + if one < two { + return one + } + return two +} diff --git a/ghttp/handlers_test.go b/ghttp/handlers_test.go new file mode 100644 index 000000000..9380dabd2 --- /dev/null +++ b/ghttp/handlers_test.go @@ -0,0 +1,114 @@ +package ghttp_test + +import ( + "fmt" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + . "github.com/onsi/gomega/ghttp" + "io" + "net/http" +) + +var _ = Describe("Handlers test", func() { + Describe("RespondWithMultiple", func() { + value := -1 + handlers := RespondWithMultiple( + func(resp http.ResponseWriter, req *http.Request) { value = 1 }, + func(resp http.ResponseWriter, req *http.Request) { value = 2 }, + func(resp http.ResponseWriter, req *http.Request) { value = 3 }, + func(resp http.ResponseWriter, req *http.Request) { value = 4 }, + func(resp http.ResponseWriter, req *http.Request) { value = 5 }, + ) + DescribeTable("Should rotate through each handler and repeat last handler", + func(expected int) { + handlers(nil, nil) + Expect(value).To(Equal(expected)) + }, + Entry("call 1", 1), + Entry("call 2", 2), + Entry("call 3", 3), + Entry("call 4", 4), + Entry("call 5", 5), + Entry("call 6", 5), + ) + }) + + Describe("RoundRobinWithMultiple", func() { + value := -1 + handlers := RoundRobinWithMultiple( + func(resp http.ResponseWriter, req *http.Request) { value = 1 }, + func(resp http.ResponseWriter, req *http.Request) { value = 2 }, + func(resp http.ResponseWriter, req *http.Request) { value = 3 }, + func(resp http.ResponseWriter, req *http.Request) { value = 4 }, + func(resp http.ResponseWriter, req *http.Request) { value = 5 }, + ) + DescribeTable("Should rotate through each and start at the beginning again", + func(expected int) { + handlers(nil, nil) + Expect(value).To(Equal(expected)) + }, + Entry("call 1", 1), + Entry("call 2", 2), + Entry("call 3", 3), + Entry("call 4", 4), + Entry("call 5", 5), + Entry("call 6", 1), + ) + }) +}) + +func ExampleRespondWithMultiple() { + server := NewServer() + server.RouteToHandler(http.MethodGet, "/example", + RespondWithMultiple( + RespondWith(http.StatusOK, "1"), + RespondWith(http.StatusOK, "2"), + RespondWith(http.StatusNotFound, "Not found"), + RespondWith(http.StatusOK, "3"), + RespondWith(http.StatusInternalServerError, "Internal server error"), + )) + for i := 0; i < 6; i++ { + resp, err := http.Get(server.URL() + "/example") + if err != nil { + panic(err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + fmt.Printf("%d - %s\n", resp.StatusCode, string(body)) + } + + // Output: 200 - 1 + // 200 - 2 + // 404 - Not found + // 200 - 3 + // 500 - Internal server error + // 500 - Internal server error +} + +func ExampleRoundRobinWithMultiple() { + server := NewServer() + server.RouteToHandler(http.MethodGet, "/example", + RoundRobinWithMultiple( + RespondWith(http.StatusOK, "1"), + RespondWith(http.StatusOK, "2"), + RespondWith(http.StatusNotFound, "Not found"), + RespondWith(http.StatusOK, "3"), + RespondWith(http.StatusInternalServerError, "Internal server error"), + )) + for i := 0; i < 6; i++ { + resp, err := http.Get(server.URL() + "/example") + if err != nil { + panic(err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + fmt.Printf("%d - %s\n", resp.StatusCode, string(body)) + } + + // Output: 200 - 1 + // 200 - 2 + // 404 - Not found + // 200 - 3 + // 500 - Internal server error + // 200 - 1 +}