feat: tests for tree and RemoveNode

pkg/mux/screen/tree/group.go

@@ -34,6 +34,20 @@ func (g *Group) addNode(node Node) {
 	g.Unlock()
 }
 
+func (g *Group) removeNode(node Node) {
+	newChildren := make([]Node, 0)
+
+	g.Lock()
+	for _, child := range g.children {
+		if child == node {
+			continue
+		}
+		newChildren = append(newChildren, child)
+	}
+	g.children = newChildren
+	g.Unlock()
+}
+
 func (g *Group) NewPane(
 	ctx context.Context,
 	stream mux.Stream,

pkg/mux/screen/tree/tree.go

@@ -83,6 +83,43 @@ func (t *Tree) NodeById(id NodeID) (Node, bool) {
 	return node, ok
 }
 
+func (t *Tree) RemoveNode(id NodeID) error {
+	if t.root.Id() == id {
+		return fmt.Errorf("cannot remove root node")
+	}
+
+	node, ok := t.NodeById(id)
+	if !ok {
+		return fmt.Errorf("node %d does not exist", id)
+	}
+
+	path := t.PathTo(node)
+	if len(path) == 0 {
+		return nil
+	}
+
+	parent := path[len(path)-2]
+	parent.(*Group).removeNode(node)
+
+	switch node := node.(type) {
+	case *Pane:
+		t.Lock()
+		node.Cancel()
+		delete(t.nodes, id)
+		t.Unlock()
+	case *Group:
+		t.Lock()
+		delete(t.nodes, id)
+		t.Unlock()
+
+		for _, child := range node.children {
+			t.RemoveNode(child.Id())
+		}
+	}
+
+	return nil
+}
+
 func (t *Tree) PaneById(id NodeID) (*Pane, bool) {
 	t.RLock()
 	defer t.RUnlock()

pkg/mux/screen/tree/tree_test.go

@@ -0,0 +1,64 @@
+package tree
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cfoust/cy/pkg/geom"
+	"github.com/cfoust/cy/pkg/mux/stream"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRoot(t *testing.T) {
+	tree := NewTree()
+	require.Equal(t, NodeID(1), tree.Root().Id())
+	require.Equal(t, 0, len(tree.Leaves()))
+}
+
+func emptyPane(g *Group) *Pane {
+	return g.NewPane(
+		context.Background(),
+		stream.NewReader(),
+		geom.DEFAULT_SIZE,
+	)
+}
+
+func TestLeaves(t *testing.T) {
+	tree := NewTree()
+	g := tree.Root().NewGroup()
+	for i := 0; i < 3; i++ {
+		emptyPane(g)
+	}
+	require.Equal(t, 3, len(g.Children()))
+	require.Equal(t, 3, len(tree.Leaves()))
+}
+
+func TestRemoveGroup(t *testing.T) {
+	tree := NewTree()
+	g := tree.Root().NewGroup()
+	for i := 0; i < 3; i++ {
+		emptyPane(g)
+	}
+
+	tree.RemoveNode(g.Id())
+	require.Equal(t, 0, len(tree.Leaves()))
+}
+
+func TestRemoveNode(t *testing.T) {
+	tree := NewTree()
+	g := tree.Root().NewGroup()
+	for i := 0; i < 3; i++ {
+		emptyPane(g)
+	}
+
+	child := g.Children()[1]
+	tree.RemoveNode(child.Id())
+	require.Equal(t, 2, len(tree.Leaves()))
+	require.Equal(t, 2, len(g.Children()))
+}
+
+func TestRemoveRoot(t *testing.T) {
+	tree := NewTree()
+	require.Error(t, tree.RemoveNode(tree.Root().Id()))
+}

pkg/mux/screen/tree/utils.go

@@ -38,7 +38,6 @@ func getLeaves(node Node) (result []Node) {
 				result,
 				getLeaves(child)...,
 			)
-			return
 		}
 	}