|
|
|
@ -17,6 +17,7 @@ type Transport struct { |
|
|
|
|
kind transport.Kind |
|
|
|
|
endpoint string |
|
|
|
|
operation string |
|
|
|
|
headers *mockHeader |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (tr *Transport) Kind() transport.Kind { |
|
|
|
@ -32,13 +33,33 @@ func (tr *Transport) Operation() string { |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (tr *Transport) RequestHeader() transport.Header { |
|
|
|
|
return nil |
|
|
|
|
return tr.headers |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (tr *Transport) ReplyHeader() transport.Header { |
|
|
|
|
return nil |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
type mockHeader struct { |
|
|
|
|
m map[string]string |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (m *mockHeader) Get(key string) string { |
|
|
|
|
return m.m[key] |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (m *mockHeader) Set(key, value string) { |
|
|
|
|
m.m[key] = value |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (m *mockHeader) Keys() []string { |
|
|
|
|
keys := make([]string, 0, len(m.m)) |
|
|
|
|
for k := range m.m { |
|
|
|
|
keys = append(keys, k) |
|
|
|
|
} |
|
|
|
|
return keys |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func TestMatch(t *testing.T) { |
|
|
|
|
tests := []struct { |
|
|
|
|
name string |
|
|
|
@ -139,7 +160,7 @@ func TestFunc(t *testing.T) { |
|
|
|
|
t.Log(req) |
|
|
|
|
return "reply", nil |
|
|
|
|
} |
|
|
|
|
next = Server(testMiddleware).Match(func(operation string) bool { |
|
|
|
|
next = Server(testMiddleware).Match(func(ctx context.Context, operation string) bool { |
|
|
|
|
if strings.HasPrefix(operation, "/go-kratos.dev") || strings.HasSuffix(operation, "world") { |
|
|
|
|
return true |
|
|
|
|
} |
|
|
|
@ -152,6 +173,66 @@ func TestFunc(t *testing.T) { |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func TestHeaderFunc(t *testing.T) { |
|
|
|
|
tests := []struct { |
|
|
|
|
name string |
|
|
|
|
ctx context.Context |
|
|
|
|
}{ |
|
|
|
|
{ |
|
|
|
|
name: "/hello.Update/world", |
|
|
|
|
ctx: transport.NewServerContext(context.Background(), &Transport{ |
|
|
|
|
operation: "/hello.Update/world", |
|
|
|
|
headers: &mockHeader{map[string]string{"X-Test": "test"}}, |
|
|
|
|
}), |
|
|
|
|
}, |
|
|
|
|
{ |
|
|
|
|
name: "/hi.Create/world", |
|
|
|
|
ctx: transport.NewServerContext(context.Background(), &Transport{ |
|
|
|
|
operation: "/hi.Create/world", |
|
|
|
|
headers: &mockHeader{map[string]string{"X-Test": "test2", "go-kratos": "kratos"}}, |
|
|
|
|
}), |
|
|
|
|
}, |
|
|
|
|
{ |
|
|
|
|
name: "/test.Name/1234", |
|
|
|
|
ctx: transport.NewServerContext(context.Background(), &Transport{ |
|
|
|
|
operation: "/test.Name/1234", |
|
|
|
|
headers: &mockHeader{map[string]string{"X-Test": "test3"}}, |
|
|
|
|
}), |
|
|
|
|
}, |
|
|
|
|
{ |
|
|
|
|
name: "/go-kratos.dev/kratos", |
|
|
|
|
ctx: transport.NewServerContext(context.Background(), &Transport{ |
|
|
|
|
operation: "/go-kratos.dev/kratos", |
|
|
|
|
headers: &mockHeader{map[string]string{"X-Test": "test"}}, |
|
|
|
|
}), |
|
|
|
|
}, |
|
|
|
|
} |
|
|
|
|
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(ctx context.Context, operation string) bool { |
|
|
|
|
tr, ok := transport.FromServerContext(ctx) |
|
|
|
|
if !ok { |
|
|
|
|
return false |
|
|
|
|
} |
|
|
|
|
if tr.RequestHeader().Get("X-Test") == "test" { |
|
|
|
|
return true |
|
|
|
|
} |
|
|
|
|
if tr.RequestHeader().Get("go-kratos") == "kratos" { |
|
|
|
|
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") |
|
|
|
|