From 0d8c776299422b8a7bba86b79b3ecc973780780f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Flc=E3=82=9B?= <four_leaf_clover@foxmail.com>
Date: Sat, 17 Feb 2024 21:57:16 +0800
Subject: [PATCH] feat(middleware/request): Add a feature about supporting
 event distribution middleware.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Flc゛ <four_leaf_clover@foxmail.com>
---
 middleware/request/request.go      | 65 +++++++++++++++++++++++++++++
 middleware/request/request_test.go | 66 ++++++++++++++++++++++++++++++
 2 files changed, 131 insertions(+)
 create mode 100644 middleware/request/request.go
 create mode 100644 middleware/request/request_test.go

diff --git a/middleware/request/request.go b/middleware/request/request.go
new file mode 100644
index 00000000..c654ed66
--- /dev/null
+++ b/middleware/request/request.go
@@ -0,0 +1,65 @@
+package request
+
+import (
+	"context"
+
+	"github.com/go-kratos/kratos/v2/middleware"
+
+	"github.com/go-kratos-ecosystem/components/v2/event"
+)
+
+const (
+	BeforeName = "request.before"
+	AfterName  = "request.after"
+)
+
+type From string
+
+const (
+	FromClient From = "client"
+	FromServer From = "server"
+)
+
+type BeforeEvent struct {
+	Ctx  context.Context
+	Req  interface{}
+	From From
+}
+
+type AfterEvent struct {
+	Ctx   context.Context
+	Req   interface{}
+	Reply interface{}
+	Err   error
+	From  From
+}
+
+func Server(d *event.Dispatcher) middleware.Middleware {
+	return newMiddleware(d, FromServer)
+}
+
+func Client(d *event.Dispatcher) middleware.Middleware {
+	return newMiddleware(d, FromClient)
+}
+
+func newMiddleware(d *event.Dispatcher, from From) middleware.Middleware {
+	return func(handler middleware.Handler) middleware.Handler {
+		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
+			if d != nil {
+				d.Dispatch(BeforeName, &BeforeEvent{
+					Ctx: ctx, Req: req, From: from,
+				})
+			}
+
+			reply, err = handler(ctx, req)
+
+			if d != nil {
+				d.Dispatch(AfterName, &AfterEvent{
+					Ctx: ctx, Req: req, Reply: reply, Err: err, From: from,
+				})
+			}
+
+			return
+		}
+	}
+}
diff --git a/middleware/request/request_test.go b/middleware/request/request_test.go
new file mode 100644
index 00000000..fa166b67
--- /dev/null
+++ b/middleware/request/request_test.go
@@ -0,0 +1,66 @@
+package request
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/go-kratos-ecosystem/components/v2/event"
+)
+
+var c = make(chan From, 4)
+
+type listener struct {
+	t *testing.T
+}
+
+var _ event.Listener = (*listener)(nil)
+
+func (l *listener) Listen() []event.Event {
+	return []event.Event{
+		BeforeName,
+		AfterName,
+	}
+}
+
+func (l *listener) Handle(event event.Event, data interface{}) {
+	switch e := data.(type) {
+	case *BeforeEvent:
+		assert.Equal(l.t, BeforeName, event.String())
+		assert.Equal(l.t, e.Req, "req")
+		assert.Equal(l.t, e.Ctx, context.Background())
+		c <- e.From
+	case *AfterEvent:
+		assert.Equal(l.t, AfterName, event.String())
+		assert.Equal(l.t, e.Req, "req")
+		assert.Equal(l.t, e.Reply, "reply")
+		assert.Nil(l.t, e.Err)
+		assert.Equal(l.t, e.Ctx, context.Background())
+		c <- e.From
+	}
+}
+
+func TestRequest(t *testing.T) {
+	var (
+		d       = event.NewDispatcher()
+		handler = func(_ context.Context, req interface{}) (reply interface{}, err error) {
+			assert.Equal(t, "req", req)
+			return "reply", nil
+		}
+	)
+
+	d.AddListener(&listener{t: t})
+
+	reply, err := Server(d)(handler)(context.Background(), "req")
+	assert.Equal(t, "reply", reply)
+	assert.Nil(t, err)
+	assert.Equal(t, FromServer, <-c)
+	assert.Equal(t, FromServer, <-c)
+
+	reply, err = Client(d)(handler)(context.Background(), "req")
+	assert.Equal(t, "reply", reply)
+	assert.Nil(t, err)
+	assert.Equal(t, FromClient, <-c)
+	assert.Equal(t, FromClient, <-c)
+}