Add CommandRouter.Group

maragudk · Oct 8, 2024 · bbb02da · bbb02da
1 parent 328bebb
commit bbb02da
Show file tree

Hide file tree

Showing 2 changed files with 169 additions and 16 deletions.
diff --git a/router.go b/router.go
@@ -5,6 +5,7 @@ type CommandRouter struct {
 	commands    map[string]Command
 	middlewares []Middleware
 	patterns    []string
+	subRouters  []*CommandRouter
 }
 
 func NewCommandRouter() *CommandRouter {
@@ -42,6 +43,12 @@ func (c *CommandRouter) Run(ctx Context) error {
 		}
 	}
 
+	for _, r := range c.subRouters {
+		if err := r.Run(ctx); err == nil {
+			return err
+		}
+	}
+
 	return ErrorNotFound
 }
 
@@ -57,16 +64,23 @@ func (c *CommandRouter) SubFunc(pattern string, cmd CommandFunc) {
 	c.Sub(pattern, cmd)
 }
 
-// Group commands with a new middleware stack.
+// Group commands with a new router.
+// The middleware from the parent router is copied to the new router.
 func (c *CommandRouter) Group(cb func(r *CommandRouter)) {
-	// TODO
+	r := NewCommandRouter()
+	r.middlewares = append(r.middlewares, c.middlewares...)
+	cb(r)
+	c.subRouters = append(c.subRouters, r)
 }
 
 type Middleware = func(next Command) Command
 
 // Use a middleware for all commands. If called on the root router, it will apply to all commands.
 // If called in a Group, it will apply to all commands in that group.
 func (c *CommandRouter) Use(middlewares ...Middleware) {
+	if len(c.commands) > 0 {
+		panic("cannot add middlewares after adding commands")
+	}
 	c.middlewares = append(c.middlewares, middlewares...)
 }
 

diff --git a/router_test.go b/router_test.go
@@ -10,6 +10,24 @@ import (
 )
 
 func TestCommandRouter_Run(t *testing.T) {
+	t.Run("errors on run if there is no root command", func(t *testing.T) {
+		r := clir.NewCommandRouter()
+
+		err := r.Run(clir.Context{})
+		is.Error(t, err, clir.ErrorNotFound)
+	})
+
+	t.Run("errors on run if there is no subcommand", func(t *testing.T) {
+		r := clir.NewCommandRouter()
+
+		err := r.Run(clir.Context{
+			Args: []string{"party"},
+		})
+		is.Error(t, err, clir.ErrorNotFound)
+	})
+}
+
+func TestCommandRouter_SubFunc(t *testing.T) {
 	t.Run("can route and run a root command", func(t *testing.T) {
 		r := clir.NewCommandRouter()
 
@@ -24,13 +42,6 @@ func TestCommandRouter_Run(t *testing.T) {
 		is.True(t, called)
 	})
 
-	t.Run("errors on run if there is no root command", func(t *testing.T) {
-		r := clir.NewCommandRouter()
-
-		err := r.Run(clir.Context{})
-		is.Error(t, err, clir.ErrorNotFound)
-	})
-
 	t.Run("can route and run a subcommand", func(t *testing.T) {
 		r := clir.NewCommandRouter()
 
@@ -47,40 +58,159 @@ func TestCommandRouter_Run(t *testing.T) {
 		is.NotError(t, err)
 		is.True(t, called)
 	})
+}
 
-	t.Run("errors on run if there is no subcommand", func(t *testing.T) {
+func TestCommandRouter_Use(t *testing.T) {
+	t.Run("can use middlewares on root and subcommands", func(t *testing.T) {
 		r := clir.NewCommandRouter()
 
+		m1 := newMiddleware(t, "m1")
+		m2 := newMiddleware(t, "m2")
+
+		r.Use(m1, m2)
+
+		r.SubFunc("", func(ctx clir.Context) error {
+			ctx.Println("root")
+			return nil
+		})
+
+		r.SubFunc("party", func(ctx clir.Context) error {
+			ctx.Println("party")
+			return nil
+		})
+
+		var b strings.Builder
 		err := r.Run(clir.Context{
+			Out: &b,
+		})
+		is.NotError(t, err)
+		is.Equal(t, "m1\nm2\nroot\n", b.String())
+
+		b.Reset()
+
+		err = r.Run(clir.Context{
 			Args: []string{"party"},
+			Out:  &b,
 		})
-		is.Error(t, err, clir.ErrorNotFound)
+		is.NotError(t, err)
+		is.Equal(t, "m1\nm2\nparty\n", b.String())
 	})
 
-	t.Run("can use middlewares on root and subcommands", func(t *testing.T) {
+	t.Run("panics if commands are already registered", func(t *testing.T) {
+		r := clir.NewCommandRouter()
+
+		r.SubFunc("", func(ctx clir.Context) error {
+			return nil
+		})
+
+		defer func() {
+			if rec := recover(); rec == nil {
+				t.FailNow()
+			}
+		}()
+
+		r.Use(newMiddleware(t, "m1"))
+	})
+}
+
+func TestCommandRouter_Group(t *testing.T) {
+	t.Run("can group commands with a new middleware stack", func(t *testing.T) {
 		r := clir.NewCommandRouter()
 
 		m1 := newMiddleware(t, "m1")
 		m2 := newMiddleware(t, "m2")
+		m3 := newMiddleware(t, "m3")
 
-		r.Use(m1, m2)
+		// Only apply the first one here
+		r.Use(m1)
 
 		r.SubFunc("", func(ctx clir.Context) error {
 			ctx.Println("root")
 			return nil
 		})
 
-		r.SubFunc("party", func(ctx clir.Context) error {
-			ctx.Println("party")
+		r.Group(func(r *clir.CommandRouter) {
+			r.Use(m2)
+
+			r.SubFunc("party", func(ctx clir.Context) error {
+				ctx.Println("party")
+				return nil
+			})
+		})
+
+		r.Group(func(r *clir.CommandRouter) {
+			r.Use(m3)
+
+			r.SubFunc("sleep", func(ctx clir.Context) error {
+				ctx.Println("sleep")
+				return nil
+			})
+		})
+
+		var b strings.Builder
+		err := r.Run(clir.Context{
+			Out: &b,
+		})
+		is.NotError(t, err)
+		is.Equal(t, "m1\nroot\n", b.String())
+
+		b.Reset()
+
+		err = r.Run(clir.Context{
+			Args: []string{"party"},
+			Out:  &b,
+		})
+		is.NotError(t, err)
+		is.Equal(t, "m1\nm2\nparty\n", b.String())
+
+		b.Reset()
+
+		err = r.Run(clir.Context{
+			Args: []string{"sleep"},
+			Out:  &b,
+		})
+		is.NotError(t, err)
+		is.Equal(t, "m1\nm3\nsleep\n", b.String())
+	})
+
+	t.Run("can nest groups", func(t *testing.T) {
+		r := clir.NewCommandRouter()
+
+		m1 := newMiddleware(t, "m1")
+		m2 := newMiddleware(t, "m2")
+		m3 := newMiddleware(t, "m3")
+
+		r.Use(m1)
+
+		r.SubFunc("", func(ctx clir.Context) error {
+			ctx.Println("root")
 			return nil
 		})
 
+		r.Group(func(r *clir.CommandRouter) {
+			r.Use(m2)
+
+			r.SubFunc("party", func(ctx clir.Context) error {
+				ctx.Println("party")
+				return nil
+			})
+
+			r.Group(func(r *clir.CommandRouter) {
+				r.Use(m3)
+
+				r.SubFunc("sleep", func(ctx clir.Context) error {
+					ctx.Println("sleep")
+					return nil
+				})
+			})
+		})
+
 		var b strings.Builder
 		err := r.Run(clir.Context{
 			Out: &b,
 		})
 		is.NotError(t, err)
-		is.Equal(t, "m1\nm2\nroot\n", b.String())
+		is.Equal(t, "m1\nroot\n", b.String())
 
 		b.Reset()
 
@@ -90,6 +220,15 @@ func TestCommandRouter_Run(t *testing.T) {
 		})
 		is.NotError(t, err)
 		is.Equal(t, "m1\nm2\nparty\n", b.String())
+
+		b.Reset()
+
+		err = r.Run(clir.Context{
+			Args: []string{"sleep"},
+			Out:  &b,
+		})
+		is.NotError(t, err)
+		is.Equal(t, "m1\nm2\nm3\nsleep\n", b.String())
 	})
 }