1
0
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:
bing liu 2024-01-05 11:01:17 +08:00 committed by GitHub
parent 4cabcaa35c
commit 21de240843
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 30 additions and 2 deletions

View File

@ -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),

View File

@ -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)
}
}