1 files changed, 76 insertions, 0 deletions
diff --git a/workhorse/internal/queueing/requests_test.go b/workhorse/internal/queueing/requests_test.go
new file mode 100644
index 00000000000..f1c52e5c6f5
--- /dev/null
+++ b/workhorse/internal/queueing/requests_test.go
@@ -0,0 +1,76 @@
+package queueing
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+)
+
+var httpHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	fmt.Fprintln(w, "OK")
+})
+
+func pausedHttpHandler(pauseCh chan struct{}) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		<-pauseCh
+		fmt.Fprintln(w, "OK")
+	})
+}
+
+func TestNormalRequestProcessing(t *testing.T) {
+	w := httptest.NewRecorder()
+	h := QueueRequests("Normal request processing", httpHandler, 1, 1, time.Second)
+	h.ServeHTTP(w, nil)
+	if w.Code != 200 {
+		t.Fatal("QueueRequests should process request")
+	}
+}
+
+// testSlowRequestProcessing creates a new queue,
+// then it runs a number of requests that are going through queue,
+// we return the response of first finished request,
+// where status of request can be 200, 429 or 503
+func testSlowRequestProcessing(name string, count int, limit, queueLimit uint, queueTimeout time.Duration) *httptest.ResponseRecorder {
+	pauseCh := make(chan struct{})
+	defer close(pauseCh)
+
+	handler := QueueRequests("Slow request processing: "+name, pausedHttpHandler(pauseCh), limit, queueLimit, queueTimeout)
+
+	respCh := make(chan *httptest.ResponseRecorder, count)
+
+	// queue requests to use up the queue
+	for i := 0; i < count; i++ {
+		go func() {
+			w := httptest.NewRecorder()
+			handler.ServeHTTP(w, nil)
+			respCh <- w
+		}()
+	}
+
+	// dequeue first request
+	return <-respCh
+}
+
+// TestQueueingTimeout performs 2 requests
+// the queue limit and length is 1,
+// the second request gets timed-out
+func TestQueueingTimeout(t *testing.T) {
+	w := testSlowRequestProcessing("timeout", 2, 1, 1, time.Microsecond)
+
+	if w.Code != 503 {
+		t.Fatal("QueueRequests should timeout queued request")
+	}
+}
+
+// TestQueueingTooManyRequests performs 3 requests
+// the queue limit and length is 1,
+// so the third request has to be rejected with 429
+func TestQueueingTooManyRequests(t *testing.T) {
+	w := testSlowRequestProcessing("too many requests", 3, 1, 1, time.Minute)
+
+	if w.Code != 429 {
+		t.Fatal("QueueRequests should return immediately and return too many requests")
+	}
+}