mirror of
https://github.com/go-kratos/kratos.git
synced 2025-03-17 21:07:54 +02:00
feat: add server option for NotFoundHandler and MethodNotAllowedHandler (#3131)
Co-authored-by: Miles Liu <milesliu@birentech.com>
This commit is contained in:
parent
4cabcaa35c
commit
21de240843
@ -141,6 +141,18 @@ func PathPrefix(prefix string) ServerOption {
|
||||
}
|
||||
}
|
||||
|
||||
func NotFoundHandler(handler http.Handler) ServerOption {
|
||||
return func(s *Server) {
|
||||
s.router.NotFoundHandler = handler
|
||||
}
|
||||
}
|
||||
|
||||
func MethodNotAllowedHandler(handler http.Handler) ServerOption {
|
||||
return func(s *Server) {
|
||||
s.router.MethodNotAllowedHandler = handler
|
||||
}
|
||||
}
|
||||
|
||||
// Server is an HTTP server wrapper.
|
||||
type Server struct {
|
||||
*http.Server
|
||||
@ -177,12 +189,12 @@ func NewServer(opts ...ServerOption) *Server {
|
||||
strictSlash: true,
|
||||
router: mux.NewRouter(),
|
||||
}
|
||||
srv.router.NotFoundHandler = http.DefaultServeMux
|
||||
srv.router.MethodNotAllowedHandler = http.DefaultServeMux
|
||||
for _, o := range opts {
|
||||
o(srv)
|
||||
}
|
||||
srv.router.StrictSlash(srv.strictSlash)
|
||||
srv.router.NotFoundHandler = http.DefaultServeMux
|
||||
srv.router.MethodNotAllowedHandler = http.DefaultServeMux
|
||||
srv.router.Use(srv.filter())
|
||||
srv.Server = &http.Server{
|
||||
Handler: FilterChain(srv.filters...)(srv.router),
|
||||
|
@ -371,3 +371,19 @@ func TestListener(t *testing.T) {
|
||||
t.Errorf("expected not empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotFoundHandler(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
srv := NewServer(NotFoundHandler(mux))
|
||||
if !reflect.DeepEqual(srv.router.NotFoundHandler, mux) {
|
||||
t.Errorf("expected %v got %v", mux, srv.router.NotFoundHandler)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMethodNotAllowedHandler(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
srv := NewServer(MethodNotAllowedHandler(mux))
|
||||
if !reflect.DeepEqual(srv.router.MethodNotAllowedHandler, mux) {
|
||||
t.Errorf("expected %v got %v", mux, srv.router.MethodNotAllowedHandler)
|
||||
}
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user