From 61b4890bc60a9973d8ed30e1a3bf2ea736a6c0d2 Mon Sep 17 00:00:00 2001
From: Ivan Stasiuk <brokeyourbike@gmail.com>
Date: Sat, 17 Aug 2024 19:04:49 +0200
Subject: [PATCH] feat: q middleware

---
 middleware/q.go      | 51 ++++++++++++++++++++++++++
 middleware/q_test.go | 87 ++++++++++++++++++++++++++++++++++++++++++++
 q/qrouter/router.go  |  2 +-
 3 files changed, 139 insertions(+), 1 deletion(-)
 create mode 100644 middleware/q.go
 create mode 100644 middleware/q_test.go

diff --git a/middleware/q.go b/middleware/q.go
new file mode 100644
index 0000000..ef75be5
--- /dev/null
+++ b/middleware/q.go
@@ -0,0 +1,51 @@
+package middleware
+
+import (
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"github.com/glocurrency/commons/instrumentation"
+	"github.com/glocurrency/commons/q"
+)
+
+const QMessageCtxKey = "qMessageCtx"
+
+type qctx struct {
+	locker Locker
+}
+
+func NewQCtx(locker Locker) *qctx {
+	return &qctx{locker: locker}
+}
+
+func (m *qctx) RequireValidMessage() gin.HandlerFunc {
+	return func(ctx *gin.Context) {
+		msg, err := q.NewQMessage(ctx.Request)
+		if err != nil {
+			instrumentation.NoticeError(ctx, err, "cannot parse queue message")
+			ctx.AbortWithStatus(http.StatusBadRequest)
+			return
+		}
+
+		instrumentation.NoticeInfo(ctx, "message from queue received",
+			instrumentation.WithField("queue_message", msg))
+
+		if msg.UniqueKey != "" {
+			if err := m.locker.TryToLock(ctx, msg.UniqueKey); err != nil {
+				instrumentation.NoticeError(ctx, err, "cannot lock task",
+					instrumentation.WithField("unique_key", msg.UniqueKey))
+				ctx.AbortWithStatus(http.StatusUnprocessableEntity)
+				return
+			}
+		}
+
+		ctx.Set(QMessageCtxKey, msg)
+		ctx.Next()
+	}
+}
+
+// MustGetQMessageFromContext returns the Q message from the context.
+func MustGetQMessageFromContext(ctx *gin.Context) q.QMessage {
+	msg := ctx.MustGet(QMessageCtxKey).(q.QMessage)
+	return msg
+}
diff --git a/middleware/q_test.go b/middleware/q_test.go
new file mode 100644
index 0000000..54480d8
--- /dev/null
+++ b/middleware/q_test.go
@@ -0,0 +1,87 @@
+package middleware_test
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/glocurrency/commons/middleware"
+	"github.com/glocurrency/commons/router"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestQCtx(t *testing.T) {
+	tests := []struct {
+		name       string
+		body       []byte
+		wantStatus int
+		setupMock  func(lockerMock *middleware.MockLocker)
+	}{
+		{
+			"success without lock",
+			validPubSubMsg,
+			http.StatusOK,
+			nil,
+		},
+		{
+			"success with lock",
+			validPubSubUniqueMsg,
+			http.StatusOK,
+			func(lockerMock *middleware.MockLocker) {
+				lockerMock.On("TryToLock", mock.Anything, "unique-456").Return(nil)
+			},
+		},
+		{
+			"invalid message",
+			invalidPubSubMsg,
+			http.StatusBadRequest,
+			nil,
+		},
+		{
+			"cannot lock by unique the key",
+			validPubSubUniqueMsg,
+			http.StatusUnprocessableEntity,
+			func(lockerMock *middleware.MockLocker) {
+				lockerMock.On("TryToLock", mock.Anything, "unique-456").Return(errors.New("I am an error"))
+			},
+		},
+	}
+
+	for i := range tests {
+		test := tests[i]
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			tasksMock := middleware.NewMockLocker(t)
+			if test.setupMock != nil {
+				test.setupMock(tasksMock)
+			}
+
+			m := middleware.NewQCtx(tasksMock)
+			req := httptest.NewRequest(http.MethodPost, "/", nil)
+
+			if test.body != nil {
+				req.Body = io.NopCloser(bytes.NewReader(test.body))
+			}
+
+			w := httptest.NewRecorder()
+			router := router.NewRouterWithValidation()
+			router.Use(m.RequireValidMessage())
+			router.POST("/", func(ctx *gin.Context) {
+				msg := middleware.MustGetQMessageFromContext(ctx)
+				assert.NotNil(t, msg)
+
+				ctx.String(http.StatusOK, "the end.")
+			})
+			router.ServeHTTP(w, req)
+
+			require.Equal(t, test.wantStatus, w.Code)
+		})
+	}
+}
diff --git a/q/qrouter/router.go b/q/qrouter/router.go
index 84ccca9..94474df 100644
--- a/q/qrouter/router.go
+++ b/q/qrouter/router.go
@@ -34,7 +34,7 @@ func (r *router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 func RewriteForQ(req *http.Request) error {
 	msg, err := q.NewQMessage(req)
 	if err != nil {
-		return fmt.Errorf("cannot create q message: %w", err)
+		return fmt.Errorf("cannot parse q message: %w", err)
 	}
 
 	if msg.Name == "" {