fix: unary compatible middleware

cloudwego · Feb 13, 2025 · f243586 · f243586
1 parent f8c18e4
commit f243586
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 17 deletions.
diff --git a/streamx/proto/proto_test.go b/streamx/proto/proto_test.go
@@ -27,8 +27,8 @@ import (
 	"github.com/bytedance/gopkg/cloud/metainfo"
 	"github.com/cloudwego/kitex/client"
 	"github.com/cloudwego/kitex/pkg/endpoint"
-	cep "github.com/cloudwego/kitex/pkg/endpoint/client"
-	sep "github.com/cloudwego/kitex/pkg/endpoint/server"
+	"github.com/cloudwego/kitex/pkg/endpoint/cep"
+	"github.com/cloudwego/kitex/pkg/endpoint/sep"
 	"github.com/cloudwego/kitex/pkg/kerrors"
 	"github.com/cloudwego/kitex/pkg/klog"
 	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata"

diff --git a/streamx/thrift/thrift_test.go b/streamx/thrift/thrift_test.go
@@ -22,13 +22,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/cloudwego/kitex-tests/streamx"
-
 	"github.com/bytedance/gopkg/cloud/metainfo"
 	"github.com/cloudwego/kitex/client"
 	"github.com/cloudwego/kitex/pkg/endpoint"
-	cep "github.com/cloudwego/kitex/pkg/endpoint/client"
-	sep "github.com/cloudwego/kitex/pkg/endpoint/server"
+	"github.com/cloudwego/kitex/pkg/endpoint/cep"
+	"github.com/cloudwego/kitex/pkg/endpoint/sep"
 	"github.com/cloudwego/kitex/pkg/kerrors"
 	"github.com/cloudwego/kitex/pkg/klog"
 	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata"
@@ -41,6 +39,7 @@ import (
 
 	"github.com/cloudwego/kitex-tests/pkg/test"
 	"github.com/cloudwego/kitex-tests/pkg/utils/serverutils"
+	"github.com/cloudwego/kitex-tests/streamx"
 	"github.com/cloudwego/kitex-tests/streamx/kitex_gen/echo"
 	"github.com/cloudwego/kitex-tests/streamx/kitex_gen/echo/testservice"
 )

diff --git a/thrift_streaming/grpcpb_test.go b/thrift_streaming/grpcpb_test.go
@@ -238,19 +238,13 @@ func TestGRPCPBServerMiddleware(t *testing.T) {
 	svr := RunGRPCPBServer(&GRPCPBServiceImpl{}, addr,
 		server.WithMiddleware(func(e endpoint.Endpoint) endpoint.Endpoint {
 			return func(ctx context.Context, req, resp interface{}) (err error) {
-				method, _ := kitexutil.GetMethod(ctx)
 				_, ok := req.(*streaming.Args)
-				isKitexPB, _ := metainfo.GetValue(ctx, "TEST_KITEX_PB")
-				if !ok && isKitexPB != "1" {
-					return fmt.Errorf("invalid request type: %T, method = %v", req, method)
-				}
-				if isKitexPB == "1" {
-					req.(*grpcpbservice.EchoPingPongArgs).Req.Message += "_mwReq"
+				if ok {
+					return e(ctx, req, resp)
 				}
+				req.(*grpcpbservice.EchoPingPongArgs).Req.Message += "_mwReq"
 				err = e(ctx, req, resp)
-				if isKitexPB == "1" {
-					resp.(*grpcpbservice.EchoPingPongResult).Success.Message += "_mwResp"
-				}
+				resp.(*grpcpbservice.EchoPingPongResult).Success.Message += "_mwResp"
 				return
 			}
 		}),
@@ -282,7 +276,7 @@ func TestGRPCPBServerMiddleware(t *testing.T) {
 			Message: "EchoPingPong",
 		})
 		test.Assert(t, err == nil)
-		test.Assert(t, resp.Message == "EchoPingPong_recv_send", resp.Message)
+		test.Assert(t, resp.Message == "EchoPingPong_mwReq_mwResp", resp.Message)
 	})
 
 	t.Run("client", func(t *testing.T) {