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

     1  package tree
     2  
     3  import (
     4  	"bytes"
     5  
     6  	. "github.com/onsi/ginkgo/v2/dsl/core"
     7  	. "github.com/onsi/gomega"
     8  	"github.com/pyroscope-io/pyroscope/pkg/storage/dict"
     9  )
    10  
    11  var _ = Describe("truncation", func() {
    12  	defer GinkgoRecover()
    13  
    14  	BeforeEach(func() {
    15  		// we override these two to better see what's going on
    16  		lostDuringSerializationName = []byte("ls")
    17  		lostDuringRenderingName = jsonableSlice("lr")
    18  	})
    19  
    20  	AfterEach(func() {
    21  		lostDuringSerializationName = []byte("other")
    22  		lostDuringRenderingName = jsonableSlice("other")
    23  	})
    24  
    25  	Context("small tree", func() {
    26  		var treeA *Tree
    27  		// treeA := New()
    28  		BeforeEach(func() {
    29  			treeA = New()
    30  			treeA.Insert([]byte("a"), uint64(1))
    31  			treeA.Insert([]byte("b"), uint64(2))
    32  			treeA.Insert([]byte("c"), uint64(3))
    33  		})
    34  
    35  		Context("with dictionary", func() {
    36  			d := dict.New()
    37  			It("after serialization drops node 'a'", func() {
    38  				buf := &bytes.Buffer{}
    39  				treeA.SerializeTruncate(d, 3, buf)
    40  				b := buf.Bytes()
    41  				treeB, err := Deserialize(d, bytes.NewReader(b))
    42  				Expect(err).ToNot(HaveOccurred())
    43  				Expect(treeB.StringWithEmpty()).To(Equal(treeStr(`"b" 2|"c" 3|"ls" 1|`)))
    44  
    45  				treeA.Insert([]byte("d"), uint64(1))
    46  
    47  				buf = &bytes.Buffer{}
    48  				treeA.SerializeTruncate(d, 3, buf)
    49  				b = buf.Bytes()
    50  
    51  				treeC, err := Deserialize(d, bytes.NewReader(b))
    52  				Expect(err).ToNot(HaveOccurred())
    53  				Expect(treeC.StringWithEmpty()).To(Equal(treeStr(`"b" 2|"c" 3|"ls" 2|`)))
    54  				Expect(treeC.StringWithEmpty()).To(Equal(treeA.StringWithEmpty()))
    55  			})
    56  		})
    57  
    58  		Context("without dictionary", func() {
    59  			It("after serialization drops node 'a'", func() {
    60  				buf := &bytes.Buffer{}
    61  				treeA.SerializeTruncateNoDict(3, buf)
    62  				treeB, err := DeserializeNoDict(buf)
    63  				Expect(err).ToNot(HaveOccurred())
    64  				Expect(treeB.StringWithEmpty()).To(Equal(treeStr(`"b" 2|"c" 3|"ls" 1|`)))
    65  			})
    66  		})
    67  	})
    68  
    69  	// TODO(petethepig): add more tests
    70  })