feat(kitexutil): add util api for getting real request and response (c…

…loudwego#1198)
seabyte7 · Dec 8, 2023 · 4db8525 · 4db8525
1 parent fa05fa2
commit 4db8525
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 0 deletions.
diff --git a/pkg/utils/kitexutil/kitexutil.go b/pkg/utils/kitexutil/kitexutil.go
@@ -22,6 +22,7 @@ import (
 	"net"
 
 	"github.com/cloudwego/kitex/pkg/rpcinfo"
+	"github.com/cloudwego/kitex/pkg/utils"
 )
 
 // GetCaller is used to get the Service Name of the caller.
@@ -127,3 +128,21 @@ func GetRPCInfo(ctx context.Context) (rpcinfo.RPCInfo, bool) {
 	}
 	return ri, true
 }
+
+// GetRealReqFromKitexArgs assert the req to be KitexArgs and return the real request if succeeded, otherwise return nil.
+// This method should be used in the middleware.
+func GetRealReqFromKitexArgs(req interface{}) interface{} {
+	if arg, ok := req.(utils.KitexArgs); ok {
+		return arg.GetFirstArgument()
+	}
+	return nil
+}
+
+// GetRealRespFromKitexResult assert the req to be KitexResult and return the real response if succeeded, otherwise return nil.
+// This method should be used in the middleware.
+func GetRealRespFromKitexResult(resp interface{}) interface{} {
+	if res, ok := resp.(utils.KitexResult); ok {
+		return res.GetResult()
+	}
+	return nil
+}
diff --git a/pkg/utils/kitexutil/kitexutil_test.go b/pkg/utils/kitexutil/kitexutil_test.go
@@ -23,6 +23,7 @@ import (
 	"reflect"
 	"testing"
 
+	mocks "github.com/cloudwego/kitex/internal/mocks/thrift/fast"
 	"github.com/cloudwego/kitex/internal/test"
 	"github.com/cloudwego/kitex/pkg/rpcinfo"
 	"github.com/cloudwego/kitex/pkg/utils"
@@ -264,6 +265,54 @@ func TestGetCtxTransportProtocol(t *testing.T) {
 	}
 }
 
+func TestGetRealRequest(t *testing.T) {
+	req := &mocks.MockReq{}
+	arg := &mocks.MockTestArgs{Req: req}
+	type args struct {
+		req interface{}
+	}
+	tests := []struct {
+		name string
+		args args
+		want interface{}
+	}{
+		{name: "success", args: args{arg}, want: req},
+		{name: "nil input", args: args{nil}, want: nil},
+		{name: "wrong interface", args: args{req}, want: nil},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := GetRealReqFromKitexArgs(tt.args.req); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("GetRealReqFromKitexArgs() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestGetRealResponse(t *testing.T) {
+	success := "success"
+	result := &mocks.MockTestResult{Success: &success}
+	type args struct {
+		resp interface{}
+	}
+	tests := []struct {
+		name string
+		args args
+		want interface{}
+	}{
+		{name: "success", args: args{result}, want: &success},
+		{name: "nil input", args: args{nil}, want: nil},
+		{name: "wrong interface", args: args{success}, want: nil},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := GetRealRespFromKitexResult(tt.args.resp); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("GetRealRespFromKitexResult() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func buildRPCInfo() rpcinfo.RPCInfo {
 	from := rpcinfo.NewEndpointInfo(caller, fromMethod, fromAddr, nil)
 	to := rpcinfo.NewEndpointInfo(callee, method, nil, nil)