From 7024c7603299f5044ca75ac6bf6e9e1e742b5aac Mon Sep 17 00:00:00 2001
From: Muyao CHEN <muyao.chen@intersec.com>
Date: Thu, 3 Oct 2024 13:41:32 +0200
Subject: [PATCH] feat: Add request id middleware test

---
 README.md                                 |  8 ++++
 go.mod                                    |  3 ++
 internal/pkg/middleware/requestid.go      |  9 ++++
 internal/pkg/middleware/requestid_test.go | 52 +++++++++++++++++++++++
 4 files changed, 72 insertions(+)
 create mode 100644 internal/pkg/middleware/requestid.go
 create mode 100644 internal/pkg/middleware/requestid_test.go

diff --git a/README.md b/README.md
index 2ae7ba9..42315fc 100644
--- a/README.md
+++ b/README.md
@@ -116,3 +116,11 @@ but it is also output to files. The log files can then be fetched to
 #### Version
 
 Add versioning into the app.
+
+### 2024/10/03
+
+Set up the web server with some necessary/nice to have middlewares.
+
+- Recovery, Logger (already included in Default mode)
+- CORS
+- RequestId
diff --git a/go.mod b/go.mod
index 7f44e6f..0f9b00d 100644
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.19.0
+	github.com/stretchr/testify v1.9.0
 	go.uber.org/zap v1.27.0
 )
 
@@ -17,6 +18,7 @@ require (
 	github.com/bytedance/sonic/loader v0.1.1 // indirect
 	github.com/cloudwego/base64x v0.1.4 // indirect
 	github.com/cloudwego/iasm v0.2.0 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/fatih/color v1.14.1 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.3 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
@@ -37,6 +39,7 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
+	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
diff --git a/internal/pkg/middleware/requestid.go b/internal/pkg/middleware/requestid.go
new file mode 100644
index 0000000..308f7fd
--- /dev/null
+++ b/internal/pkg/middleware/requestid.go
@@ -0,0 +1,9 @@
+package middleware
+
+import "github.com/gin-gonic/gin"
+
+func RequestID() gin.HandlerFunc {
+	return func(ctx *gin.Context) {
+		ctx.Set("X-Request-Id", "123")
+	}
+}
diff --git a/internal/pkg/middleware/requestid_test.go b/internal/pkg/middleware/requestid_test.go
new file mode 100644
index 0000000..a9f0014
--- /dev/null
+++ b/internal/pkg/middleware/requestid_test.go
@@ -0,0 +1,52 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/assert"
+)
+
+const requestID = "X-Request-Id"
+
+type header struct {
+	Key   string
+	Value string
+}
+
+func performRequest(
+	r http.Handler,
+	method, path string,
+	headers ...header,
+) *httptest.ResponseRecorder {
+	req := httptest.NewRequest(method, path, nil)
+	for _, h := range headers {
+		req.Header.Add(h.Key, h.Value)
+	}
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+	return w
+}
+
+func TestRequestID(t *testing.T) {
+	r := gin.New()
+	r.Use(RequestID())
+	var got string
+	wanted := "123"
+
+	r.GET("/example", func(c *gin.Context) {
+		got = c.GetString(requestID)
+		c.Status(http.StatusOK)
+	})
+
+	r.POST("/example", func(c *gin.Context) {
+		got = c.GetString(requestID)
+		c.String(http.StatusAccepted, "ok")
+	})
+
+	// Test with Request ID
+	_ = performRequest(r, "GET", "/example?a=100", header{requestID, wanted})
+	assert.Equal(t, "123", got)
+}