diff --git a/router.go b/router.go index 1a3fb08..7942ddb 100644 --- a/router.go +++ b/router.go @@ -46,10 +46,12 @@ func NewRouter(r *router.Router, logger log.Logger, endpointOverride ...Endpoint } } -func (ro *Router) RegisterCallbackServer(srv server.CallbackServer) { +func (ro *Router) RegisterCallbackServer(srv server.CallbackServer) *Router { ro.Register(ro.CallbackEndpoint, router.NewServer(srv.Callback).SetLogger(ro.Logger)) + return ro } -func (ro *Router) RegisterTokenServer(srv server.TokenServer) { +func (ro *Router) RegisterTokenServer(srv server.TokenServer) *Router { ro.Register(ro.TokenEndpoint, router.NewServer(srv.Token).SetLogger(ro.Logger)) + return ro }