tree_test.go 1.5 KB
package utils

import (
	"github.com/stretchr/testify/assert"
	"strconv"
	"testing"
)

func Test_Tree(t *testing.T) {
	table := []struct {
		Input   []TreeNode
		Text    string
		Except  []string
		Except2 []string
	}{
		{
			Input: []TreeNode{
				&st{Id: 1, Pid: 0},
				&st{Id: 2, Pid: 1}, &st{Id: 3, Pid: 1}, &st{Id: 4, Pid: 1},
				&st{Id: 5, Pid: 3},
				&st{Id: 6, Pid: 5}, &st{Id: 7, Pid: 5}},
			Text: `
树形结构:	 
	 1
2    3    4
     5
    6 7 
`,
			Except:  []string{"5", "6", "7"},
			Except2: []string{"2", "4", "6", "7"},
		},
	}

	for i := range table {
		tree := NewTree(table[i].Input)
		out := tree.AllChildNode(&st{Id: 5, Pid: 3})
		var res []string = treeNodeResults(out)
		assert.Equal(t, res, table[i].Except)

		out = tree.AllLeafNode(nil) //tree.Root()
		res = treeNodeResults(out)
		assert.Equal(t, res, table[i].Except2)

		root := tree.Root()
		assert.Equal(t, root.ID(), "1")

		//tree.Add(&st{Id:10,Pid: 7})
		//
		//out = tree.AllLeafNode(tree.Root())
		//res  = treeNodeResults(out)
		//assert.Equal(t, res, []string{"2", "4", "6", "10"})

		out = tree.TreeNodePaths(&st{Id: 7, Pid: 5})
		res = treeNodeResults(out)
		assert.Equal(t, res, []string{"1", "3", "5", "7"})
	}
}

type st struct {
	Id  int
	Pid int
}

func (t *st) PID() string {
	return strconv.Itoa(t.Pid)
}
func (t *st) ID() string {
	return strconv.Itoa(t.Id)
}

func treeNodeResults(nodes []TreeNode) []string {
	var res []string
	for i := range nodes {
		res = append(res, nodes[i].ID())
	}
	return res
}