go/kratos
1
0
Fork 0
mirror of https://github.com/go-kratos/kratos.git synced 2025-03-17 21:07:54 +02:00
Code Issues Releases Activity

feat:middleware selector add func match (#1329)

Browse Source 
* feat:middleware selector add func match
This commit is contained in:
yuemoxi 2021-09-06 22:04:42 +08:00 committed by GitHub
parent 29a5327291
commit 8823a1c618
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 59 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
middleware/selector/selector.go
View File

@ -11,7 +11,7 @@ import (
type (
transporter func(ctx context.Context) (transport.Transporter, bool)
match func(operation string) bool
MatchFunc func(operation string) bool
)
var (
@ -32,6 +32,7 @@ type Builder struct {
prefix []string
regex []string
path []string
match MatchFunc
ms []middleware.Middleware
}
@ -64,6 +65,12 @@ func (b *Builder) Path(path ...string) *Builder {
return b
}
// Match is with Builder's match
func (b *Builder) Match(fn MatchFunc) *Builder {
b.match = fn
return b
}
// Build is Builder's Build, for example: Server().Path(m1,m2).Build()
func (b *Builder) Build() middleware.Middleware {
var transporter func(ctx context.Context) (transport.Transporter, bool)
@ -72,11 +79,11 @@ func (b *Builder) Build() middleware.Middleware {
} else {
transporter = serverTransporter
}
return selector(transporter, b.match, b.ms...)
return selector(transporter, b.matchs, b.ms...)
}
// match is match operation compliance Builder
func (b *Builder) match(operation string) bool {
// matchs is match operation compliance Builder
func (b *Builder) matchs(operation string) bool {
for _, prefix := range b.prefix {
if prefixMatch(prefix, operation) {
return true
@ -92,11 +99,15 @@ func (b *Builder) match(operation string) bool {
return true
}
}
if b.match != nil {
return b.match(operation)
}
return false
}
// selector middleware
func selector(transporter transporter, match match, ms ...middleware.Middleware) middleware.Middleware {
func selector(transporter transporter, match MatchFunc, ms ...middleware.Middleware) middleware.Middleware {
return func(handler middleware.Handler) middleware.Handler {
return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
info, ok := transporter(ctx)

43
middleware/selector/selector_test.go
View File

@ -3,10 +3,12 @@ package selector
import (
"context"
"fmt"
"strings"
"testing"
"github.com/go-kratos/kratos/v2/middleware"
"github.com/go-kratos/kratos/v2/transport"
"github.com/stretchr/testify/assert"
)
var _ transport.Transporter = &Transport{}
@ -109,6 +111,47 @@ func TestMatchClient(t *testing.T) {
}
}
func TestFunc(t *testing.T) {
tests := []struct {
name string
ctx context.Context
}{
{
name: "/hello.Update/world",
ctx: transport.NewServerContext(context.Background(), &Transport{operation: "/hello.Update/world"}),
},
{
name: "/hi.Create/world",
ctx: transport.NewServerContext(context.Background(), &Transport{operation: "/hi.Create/world"}),
},
{
name: "/test.Name/1234",
ctx: transport.NewServerContext(context.Background(), &Transport{operation: "/test.Name/1234"}),
},
{
name: "/go-kratos.dev/kratos",
ctx: transport.NewServerContext(context.Background(), &Transport{operation: "/go-kratos.dev/kratos"}),
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
next := func(ctx context.Context, req interface{}) (interface{}, error) {
t.Log(req)
return "reply", nil
}
next = Server(testMiddleware).Match(func(operation string) bool {
if strings.HasPrefix(operation, "/go-kratos.dev") || strings.HasSuffix(operation, "world") {
return true
}
return false
}).Build()(next)
reply, err := next(test.ctx, test.name)
assert.Equal(t, reply, "reply")
assert.Nil(t, err)
})
}
}
func testMiddleware(handler middleware.Handler) middleware.Handler {
return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
fmt.Println("before")