From 6aeab9bb72f9e1f50af1392cc15af126a4d19215 Mon Sep 17 00:00:00 2001
From: Kirill Sysoev <k.sysoev@me.com>
Date: Thu, 11 Apr 2024 20:14:06 +0800
Subject: [PATCH] Adds timeout middleware

---
 middleware/request/timeout.go      | 27 ++++++++++++++++++++++
 middleware/request/timeout_test.go | 36 ++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+)
 create mode 100644 middleware/request/timeout.go
 create mode 100644 middleware/request/timeout_test.go

diff --git a/middleware/request/timeout.go b/middleware/request/timeout.go
new file mode 100644
index 0000000..ea26a60
--- /dev/null
+++ b/middleware/request/timeout.go
@@ -0,0 +1,27 @@
+package request
+
+import (
+	"context"
+	"time"
+
+	"github.com/ksysoev/wasabi"
+	"github.com/ksysoev/wasabi/dispatch"
+)
+
+// NewSetTimeoutMiddleware returns a middleware that sets a timeout for each request.
+// The timeout duration is specified by the 'timeout' parameter.
+// The returned middleware is a function that takes a 'next' request handler as input
+// and returns a new request handler that applies the timeout to the incoming request.
+func NewSetTimeoutMiddleware(timeout time.Duration) func(next wasabi.RequestHandler) wasabi.RequestHandler {
+	return func(next wasabi.RequestHandler) wasabi.RequestHandler {
+		return dispatch.RequestHandlerFunc(func(conn wasabi.Connection, req wasabi.Request) error {
+			ctx := req.Context()
+			ctx, cancel := context.WithTimeout(ctx, timeout)
+			defer cancel()
+
+			req = req.WithContext(ctx)
+
+			return next.Handle(conn, req)
+		})
+	}
+}
diff --git a/middleware/request/timeout_test.go b/middleware/request/timeout_test.go
new file mode 100644
index 0000000..407c87e
--- /dev/null
+++ b/middleware/request/timeout_test.go
@@ -0,0 +1,36 @@
+package request
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/ksysoev/wasabi"
+	"github.com/ksysoev/wasabi/dispatch"
+	"github.com/ksysoev/wasabi/mocks"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestNewSetTimeoutMiddleware(t *testing.T) {
+	timeout := 5 * time.Second
+	handler := dispatch.RequestHandlerFunc(func(conn wasabi.Connection, req wasabi.Request) error {
+		// Your custom handler logic here
+		return nil
+	})
+
+	middleware := NewSetTimeoutMiddleware(timeout)(handler)
+
+	ctx := context.Background()
+
+	conn := mocks.NewMockConnection(t) // Create a mock connection
+	req := mocks.NewMockRequest(t)     // Create a mock request
+
+	req.EXPECT().Context().Return(ctx)
+	req.EXPECT().WithContext(mock.AnythingOfType("*context.timerCtx")).Return(req)
+
+	err := middleware.Handle(conn, req)
+
+	if err != nil {
+		t.Errorf("Expected no error, but got %v", err)
+	}
+}