github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/storage/tree/tree_test.go (about)

     1  package tree
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"fmt"
     7  	"strings"
     8  
     9  	. "github.com/onsi/ginkgo/v2/dsl/core"
    10  	. "github.com/onsi/gomega"
    11  	"github.com/onsi/gomega/format"
    12  )
    13  
    14  var _ = Describe("tree package", func() {
    15  	Context("Insert", func() {
    16  		tree := New()
    17  		tree.Insert([]byte("a;b"), uint64(1))
    18  		tree.Insert([]byte("a;c"), uint64(2))
    19  
    20  		It("properly sets up a tree", func() {
    21  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
    22  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(2))
    23  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
    24  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
    25  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(1)))
    26  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(2)))
    27  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(1)))
    28  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(2)))
    29  			Expect(tree.String()).To(Equal("a;b 1\na;c 2\n"))
    30  		})
    31  	})
    32  
    33  	Context("Diff", func() {
    34  		a := New()
    35  		a.Insert([]byte("a;b;c"), uint64(100))
    36  		a.Insert([]byte("a;b;c;d"), uint64(100))
    37  		a.Insert([]byte("a;b;d"), uint64(100))
    38  		a.Insert([]byte("a;e"), uint64(100))
    39  		a.Insert([]byte("a;f"), uint64(150))
    40  		a.Insert([]byte("a;h"), uint64(150))
    41  
    42  		b := New()
    43  		b.Insert([]byte("a;b;c"), uint64(120))
    44  		b.Insert([]byte("a;b;c;d"), uint64(120))
    45  		b.Insert([]byte("a;b;d"), uint64(120))
    46  		b.Insert([]byte("a;e"), uint64(100))
    47  		b.Insert([]byte("a;f"), uint64(150))
    48  		b.Insert([]byte("a;g"), uint64(20))
    49  		b.Insert([]byte("a;h"), uint64(170))
    50  
    51  		diff := a.Diff(b)
    52  		z, _ := json.MarshalIndent(diff, "", "\t")
    53  		fmt.Println(string(z))
    54  		It("properly sets up a tree", func() {
    55  			Expect(diff).To(beTree([]stack{
    56  				{"a;g", 20},
    57  				{"a;h", 20},
    58  				{"a;b;c", 20},
    59  				{"a;b;d", 20},
    60  				{"a;b;c;d", 20},
    61  			}))
    62  		})
    63  	})
    64  
    65  	Context("InsertStackString unsorted of length 1", func() {
    66  		tree := New()
    67  		tree.InsertStackString([]string{"a", "b"}, uint64(1))
    68  		tree.InsertStackString([]string{"a", "a"}, uint64(2))
    69  
    70  		It("properly sets up a tree", func() {
    71  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
    72  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(2))
    73  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
    74  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
    75  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(2)))
    76  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(1)))
    77  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(2)))
    78  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(1)))
    79  			Expect(tree.String()).To(Equal("a;a 2\na;b 1\n"))
    80  		})
    81  	})
    82  
    83  	Context("InsertStackString equal of length 1", func() {
    84  		tree := New()
    85  		tree.InsertStackString([]string{"a", "b"}, uint64(1))
    86  		tree.InsertStackString([]string{"a", "b"}, uint64(2))
    87  
    88  		It("properly sets up a tree", func() {
    89  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
    90  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(1))
    91  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
    92  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
    93  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(3)))
    94  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(3)))
    95  			Expect(tree.String()).To(Equal("a;b 3\n"))
    96  		})
    97  	})
    98  
    99  	Context("InsertStackString sorted of length 1", func() {
   100  		tree := New()
   101  		tree.InsertStackString([]string{"a", "b"}, uint64(1))
   102  		tree.InsertStackString([]string{"a", "c"}, uint64(2))
   103  
   104  		It("properly sets up a tree", func() {
   105  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
   106  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(2))
   107  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
   108  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
   109  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(1)))
   110  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(2)))
   111  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(1)))
   112  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(2)))
   113  			Expect(tree.String()).To(Equal("a;b 1\na;c 2\n"))
   114  		})
   115  	})
   116  
   117  	Context("InsertStackString sorted of different lengths", func() {
   118  		tree := New()
   119  		tree.InsertStackString([]string{"a", "b"}, uint64(1))
   120  		tree.InsertStackString([]string{"a", "ba"}, uint64(2))
   121  
   122  		It("properly sets up a tree", func() {
   123  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
   124  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(2))
   125  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
   126  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
   127  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(1)))
   128  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(2)))
   129  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(1)))
   130  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(2)))
   131  			Expect(tree.String()).To(Equal("a;b 1\na;ba 2\n"))
   132  		})
   133  	})
   134  
   135  	Context("InsertStackString unsorted of different lengths", func() {
   136  		tree := New()
   137  		tree.InsertStackString([]string{"a", "ba"}, uint64(1))
   138  		tree.InsertStackString([]string{"a", "b"}, uint64(2))
   139  
   140  		It("properly sets up a tree", func() {
   141  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
   142  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(2))
   143  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
   144  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
   145  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(2)))
   146  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(1)))
   147  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(2)))
   148  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(1)))
   149  			Expect(tree.String()).To(Equal("a;b 2\na;ba 1\n"))
   150  		})
   151  	})
   152  
   153  	Context("InsertStackString unsorted of length 2", func() {
   154  		tree := New()
   155  		tree.InsertStackString([]string{"a", "bb"}, uint64(1))
   156  		tree.InsertStackString([]string{"a", "ba"}, uint64(2))
   157  
   158  		It("properly sets up a tree", func() {
   159  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
   160  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(2))
   161  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
   162  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
   163  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(2)))
   164  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(1)))
   165  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(2)))
   166  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(1)))
   167  			Expect(tree.String()).To(Equal("a;ba 2\na;bb 1\n"))
   168  		})
   169  	})
   170  
   171  	Context("InsertStackString equal of length 2", func() {
   172  		tree := New()
   173  		tree.InsertStackString([]string{"a", "bb"}, uint64(1))
   174  		tree.InsertStackString([]string{"a", "bb"}, uint64(2))
   175  
   176  		It("properly sets up a tree", func() {
   177  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
   178  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(1))
   179  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
   180  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
   181  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(3)))
   182  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(3)))
   183  			Expect(tree.String()).To(Equal("a;bb 3\n"))
   184  		})
   185  	})
   186  
   187  	Context("InsertStackString sorted of length 2", func() {
   188  		tree := New()
   189  		tree.InsertStackString([]string{"a", "bb"}, uint64(1))
   190  		tree.InsertStackString([]string{"a", "bc"}, uint64(2))
   191  
   192  		It("properly sets up a tree", func() {
   193  			Expect(tree.root.ChildrenNodes).To(HaveLen(1))
   194  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(2))
   195  			Expect(tree.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
   196  			Expect(tree.root.ChildrenNodes[0].Total).To(Equal(uint64(3)))
   197  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(1)))
   198  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(2)))
   199  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(1)))
   200  			Expect(tree.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(2)))
   201  			Expect(tree.String()).To(Equal("a;bb 1\na;bc 2\n"))
   202  		})
   203  	})
   204  
   205  	Context("Merge", func() {
   206  		Context("similar trees", func() {
   207  			treeA := New()
   208  			treeA.Insert([]byte("a;b"), uint64(1))
   209  			treeA.Insert([]byte("a;c"), uint64(2))
   210  			It("properly sets up tree A", func() {
   211  				Expect(treeA.String()).To(Equal(treeStr(`a;b 1|a;c 2|`)))
   212  			})
   213  
   214  			treeB := New()
   215  			treeB.Insert([]byte("a;b"), uint64(4))
   216  			treeB.Insert([]byte("a;c"), uint64(8))
   217  			It("properly sets up tree B", func() {
   218  				Expect(treeB.String()).To(Equal(treeStr(`a;b 4|a;c 8|`)))
   219  			})
   220  
   221  			It("properly merges", func() {
   222  				treeA.Merge(treeB)
   223  
   224  				Expect(treeA.root.ChildrenNodes).To(HaveLen(1))
   225  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(2))
   226  				Expect(treeA.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
   227  				Expect(treeA.root.ChildrenNodes[0].Total).To(Equal(uint64(15)))
   228  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(5)))
   229  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(10)))
   230  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(5)))
   231  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(10)))
   232  				Expect(treeA.String()).To(Equal(treeStr(`a;b 5|a;c 10|`)))
   233  			})
   234  		})
   235  
   236  		Context("tree with an extra node", func() {
   237  			treeA := New()
   238  			treeA.Insert([]byte("a;b"), uint64(1))
   239  			treeA.Insert([]byte("a;c"), uint64(2))
   240  			treeA.Insert([]byte("a;e"), uint64(3))
   241  			It("properly sets up tree A", func() {
   242  				Expect(treeA.String()).To(Equal(treeStr(`a;b 1|a;c 2|a;e 3|`)))
   243  			})
   244  
   245  			treeB := New()
   246  			treeB.Insert([]byte("a;b"), uint64(4))
   247  			treeB.Insert([]byte("a;d"), uint64(8))
   248  			treeB.Insert([]byte("a;e"), uint64(12))
   249  			It("properly sets up tree B", func() {
   250  				Expect(treeB.String()).To(Equal(treeStr(`a;b 4|a;d 8|a;e 12|`)))
   251  			})
   252  
   253  			It("properly merges", func() {
   254  				treeA.Merge(treeB)
   255  
   256  				Expect(treeA.root.ChildrenNodes).To(HaveLen(1))
   257  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes).To(HaveLen(4))
   258  				Expect(treeA.root.ChildrenNodes[0].Self).To(Equal(uint64(0)))
   259  				Expect(treeA.root.ChildrenNodes[0].Total).To(Equal(uint64(30)))
   260  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[0].Self).To(Equal(uint64(5)))
   261  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[1].Self).To(Equal(uint64(2)))
   262  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[2].Self).To(Equal(uint64(8)))
   263  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[3].Self).To(Equal(uint64(15)))
   264  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[0].Total).To(Equal(uint64(5)))
   265  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[1].Total).To(Equal(uint64(2)))
   266  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[2].Total).To(Equal(uint64(8)))
   267  				Expect(treeA.root.ChildrenNodes[0].ChildrenNodes[3].Total).To(Equal(uint64(15)))
   268  				Expect(treeA.String()).To(Equal(treeStr(`a;b 5|a;c 2|a;d 8|a;e 15|`)))
   269  			})
   270  		})
   271  
   272  		Context("tree.scale", func() {
   273  			treeA := New()
   274  			treeA.Insert([]byte("a;b"), uint64(1))
   275  			treeA.Insert([]byte("a;c"), uint64(2))
   276  			treeA.Insert([]byte("a;e"), uint64(3))
   277  			treeA.Insert([]byte("a"), uint64(4))
   278  			treeA.Scale(3)
   279  			It("", func() {
   280  				Expect(treeA.String()).To(Equal(treeStr(`a 12|a;b 3|a;c 6|a;e 9|`)))
   281  			})
   282  		})
   283  	})
   284  })
   285  
   286  func treeStr(s string) string {
   287  	return strings.ReplaceAll(s, "|", "\n")
   288  }
   289  
   290  var _ = Describe("prepend", func() {
   291  	Context("prependTreeNode)", func() {
   292  		It("prepend elem", func() {
   293  			A, B, C, X := &treeNode{}, &treeNode{}, &treeNode{}, &treeNode{}
   294  			s := []*treeNode{A, B, C}
   295  			s = prependTreeNode(s, X)
   296  			Expect(s).To(HaveLen(4))
   297  			Expect(s[0]).To(Equal(X))
   298  			Expect(s[1]).To(Equal(A))
   299  			Expect(s[2]).To(Equal(B))
   300  			Expect(s[3]).To(Equal(C))
   301  		})
   302  	})
   303  	Context("prependBytes", func() {
   304  		It("prepend elem", func() {
   305  			A, B, C, X := []byte("A"), []byte("B"), []byte("C"), []byte("X")
   306  			s := [][]byte{A, B, C}
   307  			s = prependBytes(s, X)
   308  
   309  			out := bytes.Join(s, []byte(","))
   310  			Expect(string(out)).To(Equal("X,A,B,C"))
   311  		})
   312  	})
   313  })
   314  
   315  type BeTreeMatcher struct {
   316  	Expected string
   317  }
   318  
   319  type stack struct {
   320  	Name  string
   321  	Value int
   322  }
   323  
   324  func beTree(stacks []stack) *BeTreeMatcher {
   325  	var b strings.Builder
   326  	for _, s := range stacks {
   327  		_, _ = fmt.Fprintf(&b, "%s %d\n", s.Name, s.Value)
   328  	}
   329  	return &BeTreeMatcher{Expected: b.String()}
   330  }
   331  
   332  func (m *BeTreeMatcher) Match(actual interface{}) (success bool, err error) {
   333  	t, ok := actual.(*Tree)
   334  	if !ok {
   335  		return false, nil
   336  	}
   337  	return t.String() == m.Expected, nil
   338  }
   339  
   340  func (m *BeTreeMatcher) FailureMessage(actual interface{}) string {
   341  	return format.Message(actual.(*Tree).String(), "to be", m.Expected)
   342  }
   343  
   344  func (m *BeTreeMatcher) NegatedFailureMessage(actual interface{}) string {
   345  	return format.Message(actual.(*Tree).String(), "not to be", m.Expected)
   346  }