From e0590aebf4b824fd262ab2d2186c60db3ee62a32 Mon Sep 17 00:00:00 2001
From: Caleb Foust <cfoust@sqweebloid.com>
Date: Wed, 20 Sep 2023 18:43:52 -0400
Subject: [PATCH] feat: tests for tree and RemoveNode

---
 pkg/mux/screen/tree/group.go     | 14 +++++++
 pkg/mux/screen/tree/tree.go      | 37 ++++++++++++++++++
 pkg/mux/screen/tree/tree_test.go | 64 ++++++++++++++++++++++++++++++++
 pkg/mux/screen/tree/utils.go     |  1 -
 4 files changed, 115 insertions(+), 1 deletion(-)
 create mode 100644 pkg/mux/screen/tree/tree_test.go

diff --git a/pkg/mux/screen/tree/group.go b/pkg/mux/screen/tree/group.go
index 651dac52..174ffac1 100644
--- a/pkg/mux/screen/tree/group.go
+++ b/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,
diff --git a/pkg/mux/screen/tree/tree.go b/pkg/mux/screen/tree/tree.go
index 2c66adc6..d14dac7d 100644
--- a/pkg/mux/screen/tree/tree.go
+++ b/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()
diff --git a/pkg/mux/screen/tree/tree_test.go b/pkg/mux/screen/tree/tree_test.go
new file mode 100644
index 00000000..a185515b
--- /dev/null
+++ b/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()))
+}
diff --git a/pkg/mux/screen/tree/utils.go b/pkg/mux/screen/tree/utils.go
index 12b7c6c3..f9a414bc 100644
--- a/pkg/mux/screen/tree/utils.go
+++ b/pkg/mux/screen/tree/utils.go
@@ -38,7 +38,6 @@ func getLeaves(node Node) (result []Node) {
 				result,
 				getLeaves(child)...,
 			)
-			return
 		}
 	}